[AutoDiff] Support multiple differentiability result indices in SIL. (swiftlang#32206)

`DifferentiableFunctionInst` now stores result indices.
`SILAutoDiffIndices` now stores result indices instead of a source index.

`@differentiable` SIL function types may now have multiple differentiability
result indices and `@noDerivative` resutls.

`@differentiable` AST function types do not have `@noDerivative` results (yet),
so this functionality is not exposed to users.

Resolves TF-689 and TF-1256.

Infrastructural support for TF-983: supporting differentiation of `apply`
instructions with multiple active semantic results.

Author: dan-zheng

docs/ABI/Mangling.rst

@@ -589,7 +589,7 @@ mangled in to disambiguate.
   impl-function-type ::= type* 'I' FUNC-ATTRIBUTES '_'
   impl-function-type ::= type* generic-signature 'I' FUNC-ATTRIBUTES '_'
 
-  FUNC-ATTRIBUTES ::= PATTERN-SUBS? INVOCATION-SUBS? PSEUDO-GENERIC? CALLEE-ESCAPE? DIFFERENTIABILITY-KIND? CALLEE-CONVENTION FUNC-REPRESENTATION? COROUTINE-KIND? (PARAM-CONVENTION PARAM-DIFFERENTIABILITY?)* RESULT-CONVENTION* ('Y' PARAM-CONVENTION)* ('z' RESULT-CONVENTION)?
+  FUNC-ATTRIBUTES ::= PATTERN-SUBS? INVOCATION-SUBS? PSEUDO-GENERIC? CALLEE-ESCAPE? DIFFERENTIABILITY-KIND? CALLEE-CONVENTION FUNC-REPRESENTATION? COROUTINE-KIND? (PARAM-CONVENTION PARAM-DIFFERENTIABILITY?)* RESULT-CONVENTION* ('Y' PARAM-CONVENTION)* ('z' RESULT-CONVENTION RESULT-DIFFERENTIABILITY?)?
 
   PATTERN-SUBS ::= 's'                       // has pattern substitutions
   INVOCATION-SUB ::= 'I'                     // has invocation substitutions
@@ -634,6 +634,8 @@ mangled in to disambiguate.
   RESULT-CONVENTION ::= 'u'                  // unowned inner pointer
   RESULT-CONVENTION ::= 'a'                  // auto-released
 
+  RESULT-DIFFERENTIABILITY ::= 'w'            // @noDerivative
+
 For the most part, manglings follow the structure of formal language
 types.  However, in some cases it is more useful to encode the exact
 implementation details of a function type.

include/swift/AST/AutoDiff.h

@@ -173,18 +173,21 @@ enum class AutoDiffGeneratedDeclarationKind : uint8_t {
 };
 
 /// SIL-level automatic differentiation indices. Consists of:
-/// - Parameter indices: indices of parameters to differentiate with respect to.
-/// - Result index: index of the result to differentiate from.
+/// - The differentiability parameter indices.
+/// - The differentiability result indices.
 // TODO(TF-913): Remove `SILAutoDiffIndices` in favor of `AutoDiffConfig`.
-// `AutoDiffConfig` supports multiple result indices.
+// `AutoDiffConfig` additionally stores a derivative generic signature.
 struct SILAutoDiffIndices {
-  /// The index of the dependent result to differentiate from.
-  unsigned source;
-  /// The indices for independent parameters to differentiate with respect to.
+  /// The indices of independent parameters to differentiate with respect to.
   IndexSubset *parameters;
+  /// The indices of dependent results to differentiate from.
+  IndexSubset *results;
 
-  /*implicit*/ SILAutoDiffIndices(unsigned source, IndexSubset *parameters)
-      : source(source), parameters(parameters) {}
+  /*implicit*/ SILAutoDiffIndices(IndexSubset *parameters, IndexSubset *results)
+      : parameters(parameters), results(results) {
+    assert(parameters && "Parameter indices must be non-null");
+    assert(results && "Result indices must be non-null");
+  }
 
   bool operator==(const SILAutoDiffIndices &other) const;
 
@@ -202,7 +205,12 @@ struct SILAutoDiffIndices {
   SWIFT_DEBUG_DUMP;
 
   std::string mangle() const {
-    std::string result = "src_" + llvm::utostr(source) + "_wrt_";
+    std::string result = "src_";
+    interleave(
+        results->getIndices(),
+        [&](unsigned idx) { result += llvm::utostr(idx); },
+        [&] { result += '_'; });
+    result += "_wrt_";
     llvm::interleave(
         parameters->getIndices(),
         [&](unsigned idx) { result += llvm::utostr(idx); },

include/swift/AST/Types.h

@@ -3942,13 +3942,34 @@ inline bool isIndirectFormalResult(ResultConvention convention) {
   return convention == ResultConvention::Indirect;
 }
 
+/// The differentiability of a SIL function type result.
+enum class SILResultDifferentiability : unsigned {
+  /// Either differentiable or not applicable.
+  ///
+  /// - If the function type is not `@differentiable`, result
+  ///   differentiability is not applicable. This case is the default value.
+  /// - If the function type is `@differentiable`, the function is
+  ///   differentiable with respect to this result.
+  DifferentiableOrNotApplicable,
+
+  /// Not differentiable: a `@noDerivative` result.
+  ///
+  /// May be applied only to result of `@differentiable` function types.
+  /// The function type is not differentiable with respect to this result.
+  NotDifferentiable,
+};
+
 /// A result type and the rules for returning it.
 class SILResultInfo {
   llvm::PointerIntPair<CanType, 3, ResultConvention> TypeAndConvention;
+  SILResultDifferentiability Differentiability : 1;
+
 public:
   SILResultInfo() = default;
-  SILResultInfo(CanType type, ResultConvention conv)
-    : TypeAndConvention(type, conv) {
+  SILResultInfo(CanType type, ResultConvention conv,
+                SILResultDifferentiability differentiability =
+                    SILResultDifferentiability::DifferentiableOrNotApplicable)
+      : TypeAndConvention(type, conv), Differentiability(differentiability) {
     assert(type->isLegalSILType() && "SILResultInfo has illegal SIL type");
   }
 
@@ -3969,6 +3990,17 @@ class SILResultInfo {
   ResultConvention getConvention() const {
     return TypeAndConvention.getInt();
   }
+
+  SILResultDifferentiability getDifferentiability() const {
+    return Differentiability;
+  }
+
+  SILResultInfo
+  getWithDifferentiability(SILResultDifferentiability differentiability) const {
+    return SILResultInfo(getInterfaceType(), getConvention(),
+                         differentiability);
+  }
+
   /// The SIL storage type determines the ABI for arguments based purely on the
   /// formal result conventions. The actual SIL type for the result values may
   /// differ in canonical SIL. In particular, opaque values require indirect
@@ -4025,6 +4057,7 @@ class SILResultInfo {
 
   void profile(llvm::FoldingSetNodeID &id) {
     id.AddPointer(TypeAndConvention.getOpaqueValue());
+    id.AddInteger((unsigned)getDifferentiability());
   }
 
   SWIFT_DEBUG_DUMP;
@@ -4714,24 +4747,31 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   /// `@noDerivative` ones).
   IndexSubset *getDifferentiabilityParameterIndices();
 
+  /// Given that `this` is a `@differentiable` or `@differentiable(linear)`
+  /// function type, returns an `IndexSubset` corresponding to the
+  /// differentiability/linearity results (e.g. all results except the
+  /// `@noDerivative` ones).
+  IndexSubset *getDifferentiabilityResultIndices();
+
   /// Returns the `@differentiable` or `@differentiable(linear)` function type
-  /// for the given differentiability kind and parameter indices representing
-  /// differentiability/linearity parameters.
+  /// for the given differentiability kind and differentiability/linearity
+  /// parameter/result indices.
   CanSILFunctionType getWithDifferentiability(DifferentiabilityKind kind,
-                                              IndexSubset *parameterIndices);
+                                              IndexSubset *parameterIndices,
+                                              IndexSubset *resultIndices);
 
   /// Returns the SIL function type stripping differentiability kind and
   /// differentiability from all parameters.
   CanSILFunctionType getWithoutDifferentiability();
 
   /// Returns the type of the derivative function for the given parameter
-  /// indices, result index, derivative function kind, derivative function
+  /// indices, result indices, derivative function kind, derivative function
   /// generic signature (optional), and other auxiliary parameters.
   ///
   /// Preconditions:
   /// - Parameters corresponding to parameter indices must conform to
   ///   `Differentiable`.
-  /// - The result corresponding to the result index must conform to
+  /// - Results corresponding to result indices must conform to
   ///   `Differentiable`.
   ///
   /// Typing rules, given:
@@ -4803,14 +4843,8 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
   ///     function - this is more direct. It may be possible to implement
   ///     reabstraction thunk derivatives using "reabstraction thunks for
   ///     the original function's derivative", avoiding extra code generation.
-  ///
-  /// Caveats:
-  /// - We may support multiple result indices instead of a single result index
-  ///   eventually. At the SIL level, this enables differentiating wrt multiple
-  ///   function results. At the Swift level, this enables differentiating wrt
-  ///   multiple tuple elements for tuple-returning functions.
   CanSILFunctionType getAutoDiffDerivativeFunctionType(
-      IndexSubset *parameterIndices, unsigned resultIndex,
+      IndexSubset *parameterIndices, IndexSubset *resultIndices,
       AutoDiffDerivativeFunctionKind kind, Lowering::TypeConverter &TC,
       LookupConformanceFn lookupConformance,
       CanGenericSignature derivativeFunctionGenericSignature = nullptr,

include/swift/Demangling/TypeDecoder.h

@@ -94,22 +94,13 @@ enum class ImplParameterDifferentiability {
   NotDifferentiable
 };
 
-static inline llvm::Optional<ImplParameterDifferentiability>
-getDifferentiabilityFromString(StringRef string) {
-  if (string.empty())
-    return ImplParameterDifferentiability::DifferentiableOrNotApplicable;
-  if (string == "@noDerivative")
-    return ImplParameterDifferentiability::NotDifferentiable;
-  return None;
-}
-
 /// Describe a lowered function parameter, parameterized on the type
 /// representation.
 template <typename BuiltType>
 class ImplFunctionParam {
+  BuiltType Type;
   ImplParameterConvention Convention;
   ImplParameterDifferentiability Differentiability;
-  BuiltType Type;
 
 public:
   using ConventionType = ImplParameterConvention;
@@ -137,9 +128,18 @@ class ImplFunctionParam {
     return None;
   }
 
-  ImplFunctionParam(ImplParameterConvention convention,
-                    ImplParameterDifferentiability diffKind, BuiltType type)
-      : Convention(convention), Differentiability(diffKind), Type(type) {}
+  static llvm::Optional<DifferentiabilityType>
+  getDifferentiabilityFromString(StringRef string) {
+    if (string.empty())
+      return DifferentiabilityType::DifferentiableOrNotApplicable;
+    if (string == "@noDerivative")
+      return DifferentiabilityType::NotDifferentiable;
+    return None;
+  }
+
+  ImplFunctionParam(BuiltType type, ImplParameterConvention convention,
+                    ImplParameterDifferentiability diffKind)
+      : Type(type), Convention(convention), Differentiability(diffKind) {}
 
   ImplParameterConvention getConvention() const { return Convention; }
 
@@ -158,15 +158,22 @@ enum class ImplResultConvention {
   Autoreleased,
 };
 
+enum class ImplResultDifferentiability {
+  DifferentiableOrNotApplicable,
+  NotDifferentiable
+};
+
 /// Describe a lowered function result, parameterized on the type
 /// representation.
 template <typename BuiltType>
 class ImplFunctionResult {
-  ImplResultConvention Convention;
   BuiltType Type;
+  ImplResultConvention Convention;
+  ImplResultDifferentiability Differentiability;
 
 public:
   using ConventionType = ImplResultConvention;
+  using DifferentiabilityType = ImplResultDifferentiability;
 
   static llvm::Optional<ConventionType>
   getConventionFromString(StringRef conventionString) {
@@ -184,11 +191,27 @@ class ImplFunctionResult {
     return None;
   }
 
-  ImplFunctionResult(ImplResultConvention convention, BuiltType type)
-      : Convention(convention), Type(type) {}
+  static llvm::Optional<DifferentiabilityType>
+  getDifferentiabilityFromString(StringRef string) {
+    if (string.empty())
+      return DifferentiabilityType::DifferentiableOrNotApplicable;
+    if (string == "@noDerivative")
+      return DifferentiabilityType::NotDifferentiable;
+    return None;
+  }
+
+  ImplFunctionResult(
+      BuiltType type, ImplResultConvention convention,
+      ImplResultDifferentiability diffKind =
+          ImplResultDifferentiability::DifferentiableOrNotApplicable)
+      : Type(type), Convention(convention), Differentiability(diffKind) {}
 
   ImplResultConvention getConvention() const { return Convention; }
 
+  ImplResultDifferentiability getDifferentiability() const {
+    return Differentiability;
+  }
+
   BuiltType getType() const { return Type; }
 };
 
@@ -640,7 +663,7 @@ class TypeDecoder {
           if (decodeImplFunctionParam(child, parameters))
             return BuiltType();
         } else if (child->getKind() == NodeKind::ImplResult) {
-          if (decodeImplFunctionPart(child, results))
+          if (decodeImplFunctionParam(child, results))
             return BuiltType();
         } else if (child->getKind() == NodeKind::ImplErrorResult) {
           if (decodeImplFunctionPart(child, errorResults))
@@ -913,13 +936,13 @@ class TypeDecoder {
     if (!type)
       return true;
 
-    results.emplace_back(*convention, type);
+    results.emplace_back(type, *convention);
     return false;
   }
 
-  bool decodeImplFunctionParam(
-      Demangle::NodePointer node,
-      llvm::SmallVectorImpl<ImplFunctionParam<BuiltType>> &results) {
+  template <typename T>
+  bool decodeImplFunctionParam(Demangle::NodePointer node,
+                               llvm::SmallVectorImpl<T> &results) {
     // Children: `convention, differentiability?, type`
     if (node->getNumChildren() != 2 && node->getNumChildren() != 3)
       return true;
@@ -931,28 +954,26 @@ class TypeDecoder {
       return true;
 
     StringRef conventionString = conventionNode->getText();
-    auto convention =
-        ImplFunctionParam<BuiltType>::getConventionFromString(conventionString);
+    auto convention = T::getConventionFromString(conventionString);
     if (!convention)
       return true;
     BuiltType type = decodeMangledType(typeNode);
     if (!type)
       return true;
 
-    auto diffKind =
-        ImplParameterDifferentiability::DifferentiableOrNotApplicable;
+    auto diffKind = T::DifferentiabilityType::DifferentiableOrNotApplicable;
     if (node->getNumChildren() == 3) {
       auto diffKindNode = node->getChild(1);
       if (diffKindNode->getKind() != Node::Kind::ImplDifferentiability)
         return true;
       auto optDiffKind =
-          getDifferentiabilityFromString(diffKindNode->getText());
+          T::getDifferentiabilityFromString(diffKindNode->getText());
       if (!optDiffKind)
         return true;
       diffKind = *optDiffKind;
     }
 
-    results.emplace_back(*convention, diffKind, type);
+    results.emplace_back(type, *convention, diffKind);
     return false;
   }
 

include/swift/SIL/SILBuilder.h

@@ -2179,10 +2179,11 @@ class SILBuilder {
   //===--------------------------------------------------------------------===//
 
   DifferentiableFunctionInst *createDifferentiableFunction(
-      SILLocation Loc, IndexSubset *ParameterIndices, SILValue OriginalFunction,
+      SILLocation Loc, IndexSubset *ParameterIndices,
+      IndexSubset *ResultIndices, SILValue OriginalFunction,
       Optional<std::pair<SILValue, SILValue>> JVPAndVJPFunctions = None) {
     return insert(DifferentiableFunctionInst::create(
-        getModule(), getSILDebugLocation(Loc), ParameterIndices,
+        getModule(), getSILDebugLocation(Loc), ParameterIndices, ResultIndices,
         OriginalFunction, JVPAndVJPFunctions, hasOwnership()));
   }
 

include/swift/SIL/SILCloner.h

@@ -2853,6 +2853,7 @@ void SILCloner<ImplClass>::visitDifferentiableFunctionInst(
   recordClonedInstruction(
       Inst, getBuilder().createDifferentiableFunction(
                 getOpLocation(Inst->getLoc()), Inst->getParameterIndices(),
+                Inst->getResultIndices(),
                 getOpValue(Inst->getOriginalFunction()), derivativeFns));
 }