Add Differentiable requirements to pattern substitutions / pattern generic signature (#68777)

Add `Differentiable` requirements to pattern substitutions / pattern generic signature when calculating constrained function type. Also, add requirements for differentiable results as well.

Fixes #65487

Author: asl

include/swift/AST/AutoDiff.h

@@ -36,6 +36,7 @@ class AbstractFunctionDecl;
 class AnyFunctionType;
 class SourceFile;
 class SILFunctionType;
+class SILResultInfo;
 class TupleType;
 class VarDecl;
 
@@ -621,18 +622,28 @@ IndexSubset *getFunctionSemanticResultIndices(
 IndexSubset *getLoweredParameterIndices(IndexSubset *astParameterIndices,
                                         AnyFunctionType *functionType);
 
+/// Collects the semantic results of the given function type in
+/// `originalResults`. The semantic results are formal results followed by
+/// semantic result parameters, in type order.
+void
+getSemanticResults(SILFunctionType *functionType,
+                   IndexSubset *parameterIndices,
+                   SmallVectorImpl<SILResultInfo> &originalResults);
+
 /// "Constrained" derivative generic signatures require all differentiability
-/// parameters to conform to the `Differentiable` protocol.
+/// parameters / results to conform to the `Differentiable` protocol.
 ///
 /// "Constrained" transpose generic signatures additionally require all
 /// linearity parameters to satisfy `Self == Self.TangentVector`.
 ///
 /// Returns the "constrained" derivative/transpose generic signature given:
 /// - An original SIL function type.
 /// - Differentiability/linearity parameter indices.
+/// - Differentiability/linearity result indices.
 /// - A possibly "unconstrained" derivative/transpose generic signature.
 GenericSignature getConstrainedDerivativeGenericSignature(
-    SILFunctionType *originalFnTy, IndexSubset *diffParamIndices,
+    SILFunctionType *originalFnTy,
+    IndexSubset *diffParamIndices, IndexSubset *diffResultIndices,
     GenericSignature derivativeGenSig, LookupConformanceFn lookupConformance,
     bool isTranspose = false);
 

lib/AST/AutoDiff.cpp

@@ -297,8 +297,30 @@ autodiff::getLoweredParameterIndices(IndexSubset *parameterIndices,
                           loweredSILIndices);
 }
 
+/// Collects the semantic results of the given function type in
+/// `originalResults`. The semantic results are formal results followed by
+/// semantic result parameters, in type order.
+void
+autodiff::getSemanticResults(SILFunctionType *functionType,
+                             IndexSubset *parameterIndices,
+                             SmallVectorImpl<SILResultInfo> &originalResults) {
+  // Collect original formal results.
+  originalResults.append(functionType->getResults().begin(),
+                         functionType->getResults().end());
+
+  // Collect original semantic result parameters.
+  for (auto i : range(functionType->getNumParameters())) {
+    auto param = functionType->getParameters()[i];
+    if (!param.isAutoDiffSemanticResult())
+      continue;
+    if (param.getDifferentiability() != SILParameterDifferentiability::NotDifferentiable)
+      originalResults.emplace_back(param.getInterfaceType(), ResultConvention::Indirect);
+  }
+}
+
 GenericSignature autodiff::getConstrainedDerivativeGenericSignature(
-    SILFunctionType *originalFnTy, IndexSubset *diffParamIndices,
+    SILFunctionType *originalFnTy,
+    IndexSubset *diffParamIndices, IndexSubset *diffResultIndices,
     GenericSignature derivativeGenSig, LookupConformanceFn lookupConformance,
     bool isTranspose) {
   if (!derivativeGenSig)
@@ -308,21 +330,48 @@ GenericSignature autodiff::getConstrainedDerivativeGenericSignature(
   auto &ctx = originalFnTy->getASTContext();
   auto *diffableProto = ctx.getProtocol(KnownProtocolKind::Differentiable);
   SmallVector<Requirement, 4> requirements;
-  for (unsigned paramIdx : diffParamIndices->getIndices()) {
-    // Require differentiability parameters to conform to `Differentiable`.
-    auto paramType = originalFnTy->getParameters()[paramIdx].getInterfaceType();
-    Requirement req(RequirementKind::Conformance, paramType,
+
+  auto addRequirement = [&](CanType type) {
+    Requirement req(RequirementKind::Conformance, type,
                     diffableProto->getDeclaredInterfaceType());
     requirements.push_back(req);
     if (isTranspose) {
       // Require linearity parameters to additionally satisfy
       // `Self == Self.TangentVector`.
-      auto tanSpace = paramType->getAutoDiffTangentSpace(lookupConformance);
-      auto paramTanType = tanSpace->getCanonicalType();
-      Requirement req(RequirementKind::SameType, paramType, paramTanType);
+      auto tanSpace = type->getAutoDiffTangentSpace(lookupConformance);
+      auto tanType = tanSpace->getCanonicalType();
+      Requirement req(RequirementKind::SameType, type, tanType);
       requirements.push_back(req);
     }
+  };
+
+  // Require differentiability parameters to conform to `Differentiable`.
+  for (unsigned paramIdx : diffParamIndices->getIndices()) {
+    auto paramType = originalFnTy->getParameters()[paramIdx].getInterfaceType();
+    addRequirement(paramType);
+  }
+
+  // Require differentiability results to conform to `Differentiable`.
+  SmallVector<SILResultInfo, 2> originalResults;
+  getSemanticResults(originalFnTy, diffParamIndices, originalResults);
+  for (unsigned resultIdx : diffResultIndices->getIndices()) {
+    // Handle formal original result.
+    if (resultIdx < originalFnTy->getNumResults()) {
+      auto resultType = originalResults[resultIdx].getInterfaceType();
+      addRequirement(resultType);
+      continue;
+    }
+    // Handle original semantic result parameters.
+    // FIXME: Constraint generic yields when we will start supporting them
+    auto resultParamIndex = resultIdx - originalFnTy->getNumResults();
+    auto resultParamIt = std::next(
+      originalFnTy->getAutoDiffSemanticResultsParameters().begin(),
+      resultParamIndex);
+    auto paramIndex =
+      std::distance(originalFnTy->getParameters().begin(), &*resultParamIt);
+    addRequirement(originalFnTy->getParameters()[paramIndex].getInterfaceType());
   }
+
   return buildGenericSignature(ctx, derivativeGenSig,
                                /*addedGenericParams*/ {},
                                std::move(requirements));

lib/SIL/IR/SILFunctionType.cpp

@@ -366,27 +366,6 @@ getDifferentiabilityParameters(SILFunctionType *originalFnTy,
       diffParams.push_back(valueAndIndex.value());
 }
 
-/// Collects the semantic results of the given function type in
-/// `originalResults`. The semantic results are formal results followed by
-/// semantic result parameters, in type order.
-static void
-getSemanticResults(SILFunctionType *functionType,
-                   IndexSubset *parameterIndices,
-                   SmallVectorImpl<SILResultInfo> &originalResults) {
-  // Collect original formal results.
-  originalResults.append(functionType->getResults().begin(),
-                         functionType->getResults().end());
-
-  // Collect original semantic result parameters.
-  for (auto i : range(functionType->getNumParameters())) {
-    auto param = functionType->getParameters()[i];
-    if (!param.isAutoDiffSemanticResult())
-      continue;
-    if (param.getDifferentiability() != SILParameterDifferentiability::NotDifferentiable)
-      originalResults.emplace_back(param.getInterfaceType(), ResultConvention::Indirect);
-  }
-}
-
 static CanGenericSignature buildDifferentiableGenericSignature(CanGenericSignature sig,
                                                                CanType tanType,
                                                                CanType origTypeOfAbstraction) {
@@ -563,7 +542,7 @@ static CanSILFunctionType getAutoDiffDifferentialType(
   SmallVector<ProtocolConformanceRef, 4> substConformances;
 
   SmallVector<SILResultInfo, 2> originalResults;
-  getSemanticResults(originalFnTy, parameterIndices, originalResults);
+  autodiff::getSemanticResults(originalFnTy, parameterIndices, originalResults);
 
   SmallVector<SILParameterInfo, 4> diffParams;
   getDifferentiabilityParameters(originalFnTy, parameterIndices, diffParams);
@@ -647,7 +626,7 @@ static CanSILFunctionType getAutoDiffPullbackType(
   SmallVector<ProtocolConformanceRef, 4> substConformances;
 
   SmallVector<SILResultInfo, 2> originalResults;
-  getSemanticResults(originalFnTy, parameterIndices, originalResults);
+  autodiff::getSemanticResults(originalFnTy, parameterIndices, originalResults);
 
   // Given a type, returns its formal SIL parameter info.
   auto getTangentParameterConventionForOriginalResult =
@@ -791,9 +770,9 @@ static CanSILFunctionType getAutoDiffPullbackType(
                              llvm::makeArrayRef(substConformances));
   }
   return SILFunctionType::get(
-      GenericSignature(), SILFunctionType::ExtInfo(), SILCoroutineKind::None,
-      ParameterConvention::Direct_Guaranteed, pullbackParams, {},
-      pullbackResults, llvm::None, substitutions,
+      GenericSignature(), SILFunctionType::ExtInfo(), originalFnTy->getCoroutineKind(),
+      ParameterConvention::Direct_Guaranteed,
+      pullbackParams, {}, pullbackResults, llvm::None, substitutions,
       /*invocationSubstitutions*/ SubstitutionMap(), ctx);
 }
 
@@ -804,7 +783,7 @@ static CanSILFunctionType getAutoDiffPullbackType(
 /// - The invocation generic signature is replaced by the
 ///   `constrainedInvocationGenSig` argument.
 static SILFunctionType *getConstrainedAutoDiffOriginalFunctionType(
-    SILFunctionType *original, IndexSubset *parameterIndices,
+    SILFunctionType *original, IndexSubset *parameterIndices, IndexSubset *resultIndices,
     LookupConformanceFn lookupConformance,
     CanGenericSignature constrainedInvocationGenSig) {
   auto originalInvocationGenSig = original->getInvocationGenericSignature();
@@ -813,6 +792,25 @@ static SILFunctionType *getConstrainedAutoDiffOriginalFunctionType(
            constrainedInvocationGenSig->areAllParamsConcrete() &&
                "derivative function cannot have invocation generic signature "
                "when original function doesn't");
+    if (auto patternSig = original->getPatternGenericSignature()) {
+      auto constrainedPatternSig =
+        autodiff::getConstrainedDerivativeGenericSignature(
+          original, parameterIndices, resultIndices,
+          patternSig, lookupConformance).getCanonicalSignature();
+      auto constrainedPatternSubs =
+        SubstitutionMap::get(constrainedPatternSig,
+                             QuerySubstitutionMap{original->getPatternSubstitutions()},
+                             lookupConformance);
+      return SILFunctionType::get(GenericSignature(),
+                                  original->getExtInfo(), original->getCoroutineKind(),
+                                  original->getCalleeConvention(),
+                                  original->getParameters(), original->getYields(),
+                                  original->getResults(), original->getOptionalErrorResult(),
+                                  constrainedPatternSubs,
+                                  /*invocationSubstitutions*/ SubstitutionMap(), original->getASTContext(),
+                                  original->getWitnessMethodConformanceOrInvalid());
+    }
+
     return original;
   }
 
@@ -823,10 +821,10 @@ static SILFunctionType *getConstrainedAutoDiffOriginalFunctionType(
   if (!constrainedInvocationGenSig)
     return original;
   constrainedInvocationGenSig =
-      autodiff::getConstrainedDerivativeGenericSignature(
-          original, parameterIndices, constrainedInvocationGenSig,
-          lookupConformance)
-          .getCanonicalSignature();
+    autodiff::getConstrainedDerivativeGenericSignature(
+      original, parameterIndices, resultIndices,
+      constrainedInvocationGenSig,
+      lookupConformance).getCanonicalSignature();
 
   SmallVector<SILParameterInfo, 4> newParameters;
   newParameters.reserve(original->getNumParameters());
@@ -882,9 +880,10 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     return cachedResult;
 
   SILFunctionType *constrainedOriginalFnTy =
-      getConstrainedAutoDiffOriginalFunctionType(this, parameterIndices,
+      getConstrainedAutoDiffOriginalFunctionType(this, parameterIndices, resultIndices,
                                                  lookupConformance,
                                                  derivativeFnInvocationGenSig);
+
   // Compute closure type.
   CanSILFunctionType closureType;
   switch (kind) {
@@ -957,11 +956,14 @@ CanSILFunctionType SILFunctionType::getAutoDiffTransposeFunctionType(
     IndexSubset *parameterIndices, Lowering::TypeConverter &TC,
     LookupConformanceFn lookupConformance,
     CanGenericSignature transposeFnGenSig) {
+  auto &ctx = getASTContext();
+
   // Get the "constrained" transpose function generic signature.
   if (!transposeFnGenSig)
     transposeFnGenSig = getSubstGenericSignature();
   transposeFnGenSig = autodiff::getConstrainedDerivativeGenericSignature(
-                          this, parameterIndices, transposeFnGenSig,
+                          this, parameterIndices, IndexSubset::getDefault(ctx, 0),
+                          transposeFnGenSig,
                           lookupConformance, /*isLinear*/ true)
                           .getCanonicalSignature();
 

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -576,7 +576,8 @@ emitDerivativeFunctionReference(
                 .second->getDerivativeGenericSignature();
       auto derivativeConstrainedGenSig =
           autodiff::getConstrainedDerivativeGenericSignature(
-              originalFn->getLoweredFunctionType(), desiredParameterIndices,
+              originalFn->getLoweredFunctionType(),
+              desiredParameterIndices, desiredResultIndices,
               contextualDerivativeGenSig,
               LookUpConformanceInModule(context.getModule().getSwiftModule()));
       minimalWitness = SILDifferentiabilityWitness::createDefinition(

test/AutoDiff/compiler_crashers_fixed/issue-68777-lack-of-diff-requirement.swift

@@ -0,0 +1,30 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+// https://github.com/apple/swift/pull/68777
+// We used not to constrain the derivative type properly, 
+// the `Differentiable` requirement was missed if the function type had
+// pattern substitution resulting in assertion:
+//
+// Invalid type parameter in getReducedType()
+// Original type: Optional<τ_0_0.TangentVector>
+// Simplified term: τ_0_0.[Differentiable:TangentVector]
+// Longest valid prefix: τ_0_0
+// Prefix type: τ_0_0
+//
+// Requirement machine for <τ_0_0, τ_0_1>
+// Rewrite system: {
+// }
+// }
+// Property map: {
+// }
+// Conformance paths: {
+// }
+//
+// Note that the generic signature <τ_0_0, τ_0_1> should be
+// <τ_0_0 : Differentiable, τ_0_1 : Differentiable>, otherwise
+// there is no way to derive associated type τ_0_0.TangentVector
+
+import _Differentiation; 
+
+public struct D<T>: Differentiable {}
+extension D {@differentiable(reverse, wrt: self) mutating func m(r: @differentiable(reverse) (T?) -> T?) {}}