Add support for differentiable functions having multiple semantic results

Co-authored-by: Brad Larson <[email protected]>

Author: asl

include/swift/AST/AutoDiff.h

@@ -32,6 +32,7 @@
 
 namespace swift {
 
+class AbstractFunctionDecl;
 class AnyFunctionType;
 class SourceFile;
 class SILFunctionType;
@@ -398,9 +399,6 @@ class DerivativeFunctionTypeError
   enum class Kind {
     /// Original function type has no semantic results.
     NoSemanticResults,
-    /// Original function type has multiple semantic results.
-    // TODO(TF-1250): Support function types with multiple semantic results.
-    MultipleSemanticResults,
     /// Differentiability parmeter indices are empty.
     NoDifferentiabilityParameters,
     /// A differentiability parameter does not conform to `Differentiable`.
@@ -429,7 +427,6 @@ class DerivativeFunctionTypeError
   explicit DerivativeFunctionTypeError(AnyFunctionType *functionType, Kind kind)
       : functionType(functionType), kind(kind), value(Value()) {
     assert(kind == Kind::NoSemanticResults ||
-           kind == Kind::MultipleSemanticResults ||
            kind == Kind::NoDifferentiabilityParameters);
   };
 
@@ -579,6 +576,10 @@ void getFunctionSemanticResultTypes(
     SmallVectorImpl<AutoDiffSemanticFunctionResultType> &result,
     GenericEnvironment *genericEnv = nullptr);
 
+/// Returns the indices of all semantic results for a given function.
+IndexSubset *getAllFunctionSemanticResultIndices(
+    const AbstractFunctionDecl *AFD);
+
 /// Returns the lowered SIL parameter indices for the given AST parameter
 /// indices and `AnyfunctionType`.
 ///

include/swift/AST/DiagnosticsSema.def

@@ -3950,9 +3950,6 @@ NOTE(autodiff_attr_original_decl_not_same_type_context,none,
      (DescriptiveDeclKind))
 ERROR(autodiff_attr_original_void_result,none,
       "cannot differentiate void function %0", (DeclName))
-ERROR(autodiff_attr_original_multiple_semantic_results,none,
-      "cannot differentiate functions with both an 'inout' parameter and a "
-      "result", ())
 ERROR(autodiff_attr_result_not_differentiable,none,
       "can only differentiate functions with results that conform to "
       "'Differentiable', but %0 does not conform to 'Differentiable'", (Type))

lib/AST/AutoDiff.cpp

@@ -196,8 +196,16 @@ void autodiff::getFunctionSemanticResultTypes(
           functionType->getResult()->getAs<AnyFunctionType>()) {
     formalResultType = resultFunctionType->getResult();
   }
-  if (!formalResultType->isEqual(ctx.TheEmptyTupleType))
-    result.push_back({remap(formalResultType), /*isInout*/ false});
+  if (!formalResultType->isEqual(ctx.TheEmptyTupleType)) {
+    // Separate tuple elements into individual results.
+    if (formalResultType->is<TupleType>()) {
+      for (auto elt : formalResultType->castTo<TupleType>()->getElements()) {
+        result.push_back({remap(elt.getType()), /*isInout*/ false});
+      }
+    } else {
+      result.push_back({remap(formalResultType), /*isInout*/ false});
+    }
+  }
 
   // Collect `inout` parameters as semantic results.
   for (auto param : functionType->getParams())
@@ -211,6 +219,16 @@ void autodiff::getFunctionSemanticResultTypes(
   }
 }
 
+IndexSubset *
+autodiff::getAllFunctionSemanticResultIndices(const AbstractFunctionDecl *AFD) {
+  auto originalFn = AFD->getInterfaceType()->castTo<AnyFunctionType>();
+  SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+  autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+  auto numResults = semanticResults.size();
+  return IndexSubset::getDefault(
+    AFD->getASTContext(), numResults, /*includeAll*/ true);
+}
+
 // TODO(TF-874): Simplify this helper. See TF-874 for WIP.
 IndexSubset *
 autodiff::getLoweredParameterIndices(IndexSubset *parameterIndices,
@@ -395,9 +413,6 @@ void DerivativeFunctionTypeError::log(raw_ostream &OS) const {
   case Kind::NoSemanticResults:
     OS << "has no semantic results ('Void' result)";
     break;
-  case Kind::MultipleSemanticResults:
-    OS << "has multiple semantic results";
-    break;
   case Kind::NoDifferentiabilityParameters:
     OS << "has no differentiability parameters";
     break;

lib/AST/Type.cpp

@@ -5548,31 +5548,86 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
   getSubsetParameters(parameterIndices, diffParams,
                       /*reverseCurryLevels*/ !makeSelfParamFirst);
 
-  // Get the original semantic result type.
+  // Get the original non-inout semantic result types.
   SmallVector<AutoDiffSemanticFunctionResultType, 1> originalResults;
   autodiff::getFunctionSemanticResultTypes(this, originalResults);
   // Error if no original semantic results.
   if (originalResults.empty())
     return llvm::make_error<DerivativeFunctionTypeError>(
         this, DerivativeFunctionTypeError::Kind::NoSemanticResults);
-  // Error if multiple original semantic results.
-  // TODO(TF-1250): Support functions with multiple semantic results.
-  if (originalResults.size() > 1)
-    return llvm::make_error<DerivativeFunctionTypeError>(
-        this, DerivativeFunctionTypeError::Kind::MultipleSemanticResults);
-  auto originalResult = originalResults.front();
-  auto originalResultType = originalResult.type;
-
-  // Get the original semantic result type's `TangentVector` associated type.
-  auto resultTan =
-      originalResultType->getAutoDiffTangentSpace(lookupConformance);
-  // Error if original semantic result has no tangent space.
-  if (!resultTan) {
+  // Accumulate non-inout result tangent spaces.
+  SmallVector<Type, 1> resultTanTypes;
+  bool hasInoutResult = false;
+  for (auto i : range(originalResults.size())) {
+    auto originalResult = originalResults[i];
+    auto originalResultType = originalResult.type;
+    // Voids currently have a defined tangent vector, so ignore them.
+    if (originalResultType->isVoid())
+      continue;
+    if (originalResult.isInout) {
+      hasInoutResult = true;
+      continue;
+    }
+    // Get the original semantic result type's `TangentVector` associated type.
+    auto resultTan =
+        originalResultType->getAutoDiffTangentSpace(lookupConformance);
+    if (!resultTan)
+      continue;
+    auto resultTanType = resultTan->getType();
+    resultTanTypes.push_back(resultTanType);
+  }
+  // Append non-wrt inout result tangent spaces.
+  // This uses the logic from getSubsetParameters(), only operating over all
+  // parameter indices and looking for non-wrt indices.
+  SmallVector<AnyFunctionType *, 2> curryLevels;
+  // An inlined version of unwrapCurryLevels().
+  AnyFunctionType *fnTy = this;
+  while (fnTy != nullptr) {
+    curryLevels.push_back(fnTy);
+    fnTy = fnTy->getResult()->getAs<AnyFunctionType>();
+  }
+
+  SmallVector<unsigned, 2> curryLevelParameterIndexOffsets(curryLevels.size());
+  unsigned currentOffset = 0;
+  for (unsigned curryLevelIndex : llvm::reverse(indices(curryLevels))) {
+    curryLevelParameterIndexOffsets[curryLevelIndex] = currentOffset;
+    currentOffset += curryLevels[curryLevelIndex]->getNumParams();
+  }
+
+  if (!makeSelfParamFirst) {
+    std::reverse(curryLevels.begin(), curryLevels.end());
+    std::reverse(curryLevelParameterIndexOffsets.begin(),
+                 curryLevelParameterIndexOffsets.end());
+  }
+
+  for (unsigned curryLevelIndex : indices(curryLevels)) {
+    auto *curryLevel = curryLevels[curryLevelIndex];
+    unsigned parameterIndexOffset =
+        curryLevelParameterIndexOffsets[curryLevelIndex];
+    for (unsigned paramIndex : range(curryLevel->getNumParams())) {
+      if (parameterIndices->contains(parameterIndexOffset + paramIndex))
+        continue;
+
+      auto param = curryLevel->getParams()[paramIndex];
+      if (param.isInOut()) {
+        auto resultType = param.getPlainType();
+        if (resultType->isVoid())
+          continue;
+        auto resultTan = resultType->getAutoDiffTangentSpace(lookupConformance);
+        if (!resultTan)
+          continue;
+        auto resultTanType = resultTan->getType();
+        resultTanTypes.push_back(resultTanType);
+      }
+    }
+  }
+
+  // Error if no semantic result has a tangent space.
+  if (resultTanTypes.empty() && !hasInoutResult) {
     return llvm::make_error<DerivativeFunctionTypeError>(
         this, DerivativeFunctionTypeError::Kind::NonDifferentiableResult,
-        std::make_pair(originalResultType, /*index*/ 0));
+        std::make_pair(originalResults.front().type, /*index*/ 0));
   }
-  auto resultTanType = resultTan->getType();
 
   // Compute the result linear map function type.
   FunctionType *linearMapType;
@@ -5592,7 +5647,6 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     // - Original:     `(T0, inout T1, ...) -> Void`
     // - Differential: `(T0.Tan, inout T1.Tan, ...) -> Void`
     SmallVector<AnyFunctionType::Param, 4> differentialParams;
-    bool hasInoutDiffParameter = false;
     for (auto i : range(diffParams.size())) {
       auto diffParam = diffParams[i];
       auto paramType = diffParam.getPlainType();
@@ -5607,11 +5661,22 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
       }
       differentialParams.push_back(AnyFunctionType::Param(
           paramTan->getType(), Identifier(), diffParam.getParameterFlags()));
-      if (diffParam.isInOut())
-        hasInoutDiffParameter = true;
     }
-    auto differentialResult =
-        hasInoutDiffParameter ? Type(ctx.TheEmptyTupleType) : resultTanType;
+    Type differentialResult;
+    if (resultTanTypes.empty()) {
+      differentialResult = ctx.TheEmptyTupleType;
+    } else if (resultTanTypes.size() == 1) {
+      differentialResult = resultTanTypes.front();
+    } else {
+      SmallVector<TupleTypeElt, 2> differentialResults;
+      for (auto i : range(resultTanTypes.size())) {
+        auto resultTanType = resultTanTypes[i];
+        differentialResults.push_back(
+            TupleTypeElt(resultTanType, Identifier()));
+      }
+      differentialResult = TupleType::get(differentialResults, ctx);
+    }
+
     // FIXME: Verify ExtInfo state is correct, not working by accident.
     FunctionType::ExtInfo info;
     linearMapType =
@@ -5629,11 +5694,11 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     // - Original: `(T0, inout T1, ...) -> Void`
     // - Pullback: `(T1.Tan) -> (T0.Tan, ...)`
     //
-    // Case 3: original function has a wrt `inout` parameter.
-    // - Original: `(T0, inout T1, ...) -> Void`
-    // - Pullback: `(inout T1.Tan) -> (T0.Tan, ...)`
+    // Case 3: original function has wrt `inout` parameters.
+    // - Original: `(T0, inout T1, ...) -> R`
+    // - Pullback: `(R.Tan, inout T1.Tan) -> (T0.Tan, ...)`
     SmallVector<TupleTypeElt, 4> pullbackResults;
-    bool hasInoutDiffParameter = false;
+    SmallVector<AnyFunctionType::Param, 2> inoutParams;
     for (auto i : range(diffParams.size())) {
       auto diffParam = diffParams[i];
       auto paramType = diffParam.getPlainType();
@@ -5647,7 +5712,9 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
             std::make_pair(paramType, i));
       }
       if (diffParam.isInOut()) {
-        hasInoutDiffParameter = true;
+        if (paramType->isVoid())
+          continue;
+        inoutParams.push_back(diffParam);
         continue;
       }
       pullbackResults.emplace_back(paramTan->getType());
@@ -5660,12 +5727,27 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     } else {
       pullbackResult = TupleType::get(pullbackResults, ctx);
     }
-    auto flags = ParameterTypeFlags().withInOut(hasInoutDiffParameter);
-    auto pullbackParam =
-        AnyFunctionType::Param(resultTanType, Identifier(), flags);
+    // First accumulate non-inout results as pullback parameters.
+    SmallVector<FunctionType::Param, 2> pullbackParams;
+    for (auto i : range(resultTanTypes.size())) {
+      auto resultTanType = resultTanTypes[i];
+      auto flags = ParameterTypeFlags().withInOut(false);
+      pullbackParams.push_back(AnyFunctionType::Param(
+          resultTanType, Identifier(), flags));
+    }
+    // Then append inout parameters.
+    for (auto i : range(inoutParams.size())) {
+      auto inoutParam = inoutParams[i];
+      auto inoutParamType = inoutParam.getPlainType();
+      auto inoutParamTan =
+          inoutParamType->getAutoDiffTangentSpace(lookupConformance);
+      auto flags = ParameterTypeFlags().withInOut(true);
+      pullbackParams.push_back(AnyFunctionType::Param(
+          inoutParamTan->getType(), Identifier(), flags));
+    }
     // FIXME: Verify ExtInfo state is correct, not working by accident.
     FunctionType::ExtInfo info;
-    linearMapType = FunctionType::get({pullbackParam}, pullbackResult, info);
+    linearMapType = FunctionType::get(pullbackParams, pullbackResult, info);
     break;
   }
   }

lib/IRGen/IRGenMangler.h

@@ -57,9 +57,11 @@ class IRGenMangler : public Mangle::ASTMangler {
       AutoDiffDerivativeFunctionIdentifier *derivativeId) {
     beginManglingWithAutoDiffOriginalFunction(func);
     auto kind = Demangle::getAutoDiffFunctionKind(derivativeId->getKind());
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(func);
     AutoDiffConfig config(
         derivativeId->getParameterIndices(),
-        IndexSubset::get(func->getASTContext(), 1, {0}),
+        resultIndices,
         derivativeId->getDerivativeGenericSignature());
     appendAutoDiffFunctionParts("TJ", kind, config);
     appendOperator("Tj");
@@ -86,9 +88,11 @@ class IRGenMangler : public Mangle::ASTMangler {
       AutoDiffDerivativeFunctionIdentifier *derivativeId) {
     beginManglingWithAutoDiffOriginalFunction(func);
     auto kind = Demangle::getAutoDiffFunctionKind(derivativeId->getKind());
+    auto *resultIndices =
+        autodiff::getAllFunctionSemanticResultIndices(func);
     AutoDiffConfig config(
         derivativeId->getParameterIndices(),
-        IndexSubset::get(func->getASTContext(), 1, {0}),
+        resultIndices,
         derivativeId->getDerivativeGenericSignature());
     appendAutoDiffFunctionParts("TJ", kind, config);
     appendOperator("Tq");

lib/SIL/IR/SILDeclRef.cpp

@@ -1067,7 +1067,8 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
     auto *silParameterIndices = autodiff::getLoweredParameterIndices(
         derivativeFunctionIdentifier->getParameterIndices(),
         getDecl()->getInterfaceType()->castTo<AnyFunctionType>());
-    auto *resultIndices = IndexSubset::get(getDecl()->getASTContext(), 1, {0});
+    auto *resultIndices = autodiff::getAllFunctionSemanticResultIndices(
+        asAutoDiffOriginalFunction().getAbstractFunctionDecl());
     AutoDiffConfig silConfig(
         silParameterIndices, resultIndices,
         derivativeFunctionIdentifier->getDerivativeGenericSignature());

lib/SIL/IR/SILFunctionType.cpp

@@ -238,8 +238,6 @@ IndexSubset *SILFunctionType::getDifferentiabilityResultIndices() {
       resultIndices.push_back(resultAndIndex.index());
   // Check `inout` parameters.
   for (auto inoutParamAndIndex : enumerate(getIndirectMutatingParameters()))
-    // FIXME(TF-1305): The `getResults().empty()` condition is a hack.
-    //
     // Currently, an `inout` parameter can either be:
     // 1. Both a differentiability parameter and a differentiability result.
     // 2. `@noDerivative`: neither a differentiability parameter nor a
@@ -251,13 +249,8 @@ IndexSubset *SILFunctionType::getDifferentiabilityResultIndices() {
     //    cases, so supporting it is a non-goal.
     //
     // See TF-1305 for solution ideas. For now, `@noDerivative` `inout`
-    // parameters are not treated as differentiability results, unless the
-    // original function has no formal results, in which case all `inout`
     // parameters are treated as differentiability results.
-    if (getResults().empty() ||
-        inoutParamAndIndex.value().getDifferentiability() !=
-            SILParameterDifferentiability::NotDifferentiable)
-      resultIndices.push_back(getNumResults() + inoutParamAndIndex.index());
+    resultIndices.push_back(getNumResults() + inoutParamAndIndex.index());
   auto numSemanticResults =
       getNumResults() + getNumIndirectMutatingParameters();
   return IndexSubset::get(getASTContext(), numSemanticResults, resultIndices);
@@ -603,7 +596,7 @@ static CanSILFunctionType getAutoDiffDifferentialType(
       differentialResults.push_back({resultTanType, resultConv});
       continue;
     }
-    // Handle original `inout` parameter.
+    // Handle original `inout` parameters.
     auto inoutParamIndex = resultIndex - originalFnTy->getNumResults();
     auto inoutParamIt = std::next(
         originalFnTy->getIndirectMutatingParameters().begin(), inoutParamIndex);
@@ -745,7 +738,7 @@ static CanSILFunctionType getAutoDiffPullbackType(
       pullbackParams.push_back({resultTanType, paramConv});
       continue;
     }
-    // Handle original `inout` parameter.
+    // Handle `inout` parameters.
     auto inoutParamIndex = resultIndex - originalFnTy->getNumResults();
     auto inoutParamIt = std::next(
         originalFnTy->getIndirectMutatingParameters().begin(), inoutParamIndex);