Correctly infer result indices & types for multiple semantic inout results

Author: asl

include/swift/AST/AutoDiff.h

@@ -248,7 +248,12 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &s,
 /// an `inout` parameter type. Used in derivative function type calculation.
 struct AutoDiffSemanticFunctionResultType {
   Type type;
-  bool isInout;
+  unsigned index : 30;
+  bool isInout : 1;
+  bool isWrtParam : 1;
+
+  AutoDiffSemanticFunctionResultType(Type t, unsigned idx, bool inout, bool wrt)
+    : type(t), index(idx), isInout(inout), isWrtParam(wrt) { }
 };
 
 /// Key for caching SIL derivative function types.
@@ -569,19 +574,22 @@ namespace autodiff {
 /// `inout` parameter types.
 ///
 /// The function type may have at most two parameter lists.
-///
-/// Remaps the original semantic result using `genericEnv`, if specified.
-void getFunctionSemanticResultTypes(
-    AnyFunctionType *functionType,
-    SmallVectorImpl<AutoDiffSemanticFunctionResultType> &result,
-    GenericEnvironment *genericEnv = nullptr);
+void getFunctionSemanticResults(
+    const AnyFunctionType *functionType,
+    const IndexSubset *parameterIndices,
+    SmallVectorImpl<AutoDiffSemanticFunctionResultType> &resultTypes);
+
+/// Returns the indices of semantic results for a given function.
+IndexSubset *getFunctionSemanticResultIndices(
+    const AnyFunctionType *functionType,
+    const IndexSubset *parameterIndices);
 
-/// Returns the indices of all semantic results for a given function.
-IndexSubset *getAllFunctionSemanticResultIndices(
-    const AbstractFunctionDecl *AFD);
+IndexSubset *getFunctionSemanticResultIndices(
+    const AbstractFunctionDecl *AFD,
+    const IndexSubset *parameterIndices);
 
 /// Returns the lowered SIL parameter indices for the given AST parameter
-/// indices and `AnyfunctionType`.
+/// indices and `AnyFunctionType`.
 ///
 /// Notable lowering-related changes:
 /// - AST tuple parameter types are exploded when lowered to SIL.

lib/AST/AutoDiff.cpp

@@ -176,57 +176,89 @@ void AnyFunctionType::getSubsetParameters(
   }
 }
 
-void autodiff::getFunctionSemanticResultTypes(
-    AnyFunctionType *functionType,
-    SmallVectorImpl<AutoDiffSemanticFunctionResultType> &result,
-    GenericEnvironment *genericEnv) {
+void autodiff::getFunctionSemanticResults(
+    const AnyFunctionType *functionType,
+    const IndexSubset *parameterIndices,
+    SmallVectorImpl<AutoDiffSemanticFunctionResultType> &resultTypes) {
   auto &ctx = functionType->getASTContext();
 
-  // Remap type in `genericEnv`, if specified.
-  auto remap = [&](Type type) {
-    if (!genericEnv)
-      return type;
-    return genericEnv->mapTypeIntoContext(type);
-  };
-
   // Collect formal result type as a semantic result, unless it is
   // `Void`.
   auto formalResultType = functionType->getResult();
   if (auto *resultFunctionType =
-          functionType->getResult()->getAs<AnyFunctionType>()) {
+      functionType->getResult()->getAs<AnyFunctionType>())
     formalResultType = resultFunctionType->getResult();
-  }
+
+  unsigned resultIdx = 0;
   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});
+        resultTypes.emplace_back(elt.getType(), resultIdx++,
+                                 /*isInout*/ false, /*isWrt*/ false);
       }
     } else {
-      result.push_back({remap(formalResultType), /*isInout*/ false});
+      resultTypes.emplace_back(formalResultType, resultIdx++,
+                               /*isInout*/ false, /*isWrt*/ false);
     }
   }
 
-  // Collect `inout` parameters as semantic results.
-  for (auto param : functionType->getParams())
-    if (param.isInOut())
-      result.push_back({remap(param.getPlainType()), /*isInout*/ true});
-  if (auto *resultFunctionType =
-          functionType->getResult()->getAs<AnyFunctionType>()) {
-    for (auto param : resultFunctionType->getParams())
-      if (param.isInOut())
-        result.push_back({remap(param.getPlainType()), /*isInout*/ true});
-  }
+  bool addNonWrts = resultTypes.empty();
+
+  // Collect wrt `inout` parameters as semantic results
+  // As an extention, collect all (including non-wrt) inouts as results for
+  // functions returning void.
+  auto collectSemanticResults = [&](const AnyFunctionType *functionType,
+                                    unsigned curryOffset = 0) {
+    for (auto paramAndIndex : enumerate(functionType->getParams())) {
+      if (!paramAndIndex.value().isInOut())
+        continue;
+
+      unsigned idx = paramAndIndex.index() + curryOffset;
+      assert(idx < parameterIndices->getCapacity() &&
+             "invalid parameter index");
+      bool isWrt = parameterIndices->contains(idx);
+      if (addNonWrts || isWrt)
+        resultTypes.emplace_back(paramAndIndex.value().getPlainType(),
+                                 resultIdx, /*isInout*/ true, isWrt);
+      resultIdx += 1;
+    }
+  };
+
+  if (auto *resultFnType =
+      functionType->getResult()->getAs<AnyFunctionType>()) {
+    // Here we assume that the input is a function type with curried `Self`
+    assert(functionType->getNumParams() == 1 && "unexpected function type");
+
+    collectSemanticResults(resultFnType);
+    collectSemanticResults(functionType, resultFnType->getNumParams());
+  } else
+    collectSemanticResults(functionType);
 }
 
 IndexSubset *
-autodiff::getAllFunctionSemanticResultIndices(const AbstractFunctionDecl *AFD) {
-  auto originalFn = AFD->getInterfaceType()->castTo<AnyFunctionType>();
+autodiff::getFunctionSemanticResultIndices(const AnyFunctionType *functionType,
+                                           const IndexSubset *parameterIndices) {
+  auto &ctx = functionType->getASTContext();
+
   SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
-  autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
-  auto numResults = semanticResults.size();
-  return IndexSubset::getDefault(
-    AFD->getASTContext(), numResults, /*includeAll*/ true);
+  autodiff::getFunctionSemanticResults(functionType, parameterIndices,
+                                       semanticResults);
+  SmallVector<unsigned> resultIndices;
+  unsigned cap = 0;
+  for (const auto& result : semanticResults) {
+    resultIndices.push_back(result.index);
+    cap = std::max(cap, result.index + 1U);
+  }
+
+  return IndexSubset::get(ctx, cap, resultIndices);
+}
+
+IndexSubset *
+autodiff::getFunctionSemanticResultIndices(const AbstractFunctionDecl *AFD,
+                                           const IndexSubset *parameterIndices) {
+  return getFunctionSemanticResultIndices(AFD->getInterfaceType()->castTo<AnyFunctionType>(),
+                                          parameterIndices);
 }
 
 // TODO(TF-874): Simplify this helper. See TF-874 for WIP.

lib/AST/Type.cpp

@@ -5550,85 +5550,41 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
 
   // Get the original non-inout semantic result types.
   SmallVector<AutoDiffSemanticFunctionResultType, 1> originalResults;
-  autodiff::getFunctionSemanticResultTypes(this, originalResults);
+  autodiff::getFunctionSemanticResults(this, parameterIndices, originalResults);
   // Error if no original semantic results.
   if (originalResults.empty())
     return llvm::make_error<DerivativeFunctionTypeError>(
         this, DerivativeFunctionTypeError::Kind::NoSemanticResults);
+
   // Accumulate non-inout result tangent spaces.
-  SmallVector<Type, 1> resultTanTypes;
-  bool hasInoutResult = false;
+  SmallVector<Type, 1> resultTanTypes, inoutTanTypes;
   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.
+    // Error if a semantic result has no tangent space.
     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>(
+      return llvm::make_error<DerivativeFunctionTypeError>(
         this, DerivativeFunctionTypeError::Kind::NonDifferentiableResult,
-        std::make_pair(originalResults.front().type, /*index*/ 0));
+        std::make_pair(originalResultType, unsigned(originalResult.index)));
+
+    if (!originalResult.isInout)
+      resultTanTypes.push_back(resultTan->getType());
+    else if (originalResult.isInout && !originalResult.isWrtParam)
+      inoutTanTypes.push_back(resultTan->getType());
   }
 
+  // Treat non-wrt inouts as semantic results for functions returning Void
+  if (resultTanTypes.empty())
+    resultTanTypes = inoutTanTypes;
+  
   // Compute the result linear map function type.
   FunctionType *linearMapType;
   switch (kind) {
@@ -5641,24 +5597,24 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     //
     // Case 2: original function has a non-wrt `inout` parameter.
     // - Original:      `(T0, inout T1, ...) -> Void`
-    // - Differential: `(T0.Tan, ...) -> T1.Tan`
+    // - Differential:  `(T0.Tan, ...) -> T1.Tan`
     //
     // Case 3: original function has a wrt `inout` parameter.
-    // - Original:     `(T0, inout T1, ...) -> Void`
-    // - Differential: `(T0.Tan, inout T1.Tan, ...) -> Void`
+    // - Original:      `(T0, inout T1, ...) -> Void`
+    // - Differential:  `(T0.Tan, inout T1.Tan, ...) -> Void`
     SmallVector<AnyFunctionType::Param, 4> differentialParams;
     for (auto i : range(diffParams.size())) {
       auto diffParam = diffParams[i];
       auto paramType = diffParam.getPlainType();
       auto paramTan = paramType->getAutoDiffTangentSpace(lookupConformance);
       // Error if parameter has no tangent space.
-      if (!paramTan) {
+      if (!paramTan)
         return llvm::make_error<DerivativeFunctionTypeError>(
             this,
             DerivativeFunctionTypeError::Kind::
                 NonDifferentiableDifferentiabilityParameter,
             std::make_pair(paramType, i));
-      }
+
       differentialParams.push_back(AnyFunctionType::Param(
           paramTan->getType(), Identifier(), diffParam.getParameterFlags()));
     }
@@ -5704,13 +5660,13 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
       auto paramType = diffParam.getPlainType();
       auto paramTan = paramType->getAutoDiffTangentSpace(lookupConformance);
       // Error if parameter has no tangent space.
-      if (!paramTan) {
+      if (!paramTan)
         return llvm::make_error<DerivativeFunctionTypeError>(
             this,
             DerivativeFunctionTypeError::Kind::
                 NonDifferentiableDifferentiabilityParameter,
             std::make_pair(paramType, i));
-      }
+
       if (diffParam.isInOut()) {
         if (paramType->isVoid())
           continue;

lib/IRGen/IRGenMangler.h

@@ -58,7 +58,8 @@ class IRGenMangler : public Mangle::ASTMangler {
     beginManglingWithAutoDiffOriginalFunction(func);
     auto kind = Demangle::getAutoDiffFunctionKind(derivativeId->getKind());
     auto *resultIndices =
-        autodiff::getAllFunctionSemanticResultIndices(func);
+      autodiff::getFunctionSemanticResultIndices(func,
+                                                 derivativeId->getParameterIndices());
     AutoDiffConfig config(
         derivativeId->getParameterIndices(),
         resultIndices,
@@ -89,7 +90,8 @@ class IRGenMangler : public Mangle::ASTMangler {
     beginManglingWithAutoDiffOriginalFunction(func);
     auto kind = Demangle::getAutoDiffFunctionKind(derivativeId->getKind());
     auto *resultIndices =
-        autodiff::getAllFunctionSemanticResultIndices(func);
+      autodiff::getFunctionSemanticResultIndices(func,
+                                                 derivativeId->getParameterIndices());
     AutoDiffConfig config(
         derivativeId->getParameterIndices(),
         resultIndices,

lib/SIL/IR/SILDeclRef.cpp

@@ -1067,8 +1067,10 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
     auto *silParameterIndices = autodiff::getLoweredParameterIndices(
         derivativeFunctionIdentifier->getParameterIndices(),
         getDecl()->getInterfaceType()->castTo<AnyFunctionType>());
-    auto *resultIndices = autodiff::getAllFunctionSemanticResultIndices(
-        asAutoDiffOriginalFunction().getAbstractFunctionDecl());
+    // FIXME: is this correct in the presence of curried types?
+    auto *resultIndices = autodiff::getFunctionSemanticResultIndices(
+      asAutoDiffOriginalFunction().getAbstractFunctionDecl(),
+      derivativeFunctionIdentifier->getParameterIndices());
     AutoDiffConfig silConfig(
         silParameterIndices, resultIndices,
         derivativeFunctionIdentifier->getDerivativeGenericSignature());