Reworking logic for non-wrt inout parameters. Replacing single result index with all result indices in multiple places.

Author: BradLarson

lib/AST/Type.cpp

@@ -5482,11 +5482,14 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
   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;
     }
-    auto originalResultType = originalResult.type;
     // Get the original semantic result type's `TangentVector` associated type.
     auto resultTan =
         originalResultType->getAutoDiffTangentSpace(lookupConformance);
@@ -5496,19 +5499,48 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
     resultTanTypes.push_back(resultTanType);
   }
   // Append non-wrt inout result tangent spaces.
-  auto *resultFunctionType = this->getResult()->getAs<AnyFunctionType>();
-  auto sourceFunction = resultFunctionType ? resultFunctionType : this;
-  for (unsigned i : range(sourceFunction->getNumParams())) {
-    auto param = sourceFunction->getParams()[i];
-    if (parameterIndices->contains(i))
-      continue;
-    if (param.isInOut()) {
-      auto resultType = param.getPlainType();
-      auto resultTan = resultType->getAutoDiffTangentSpace(lookupConformance);
-      if (!resultTan)
+  // 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 resultTanType = resultTan->getType();
-      resultTanTypes.push_back(resultTanType);
+
+      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);
+      }
     }
   }
 
@@ -5603,6 +5635,8 @@ AnyFunctionType::getAutoDiffDerivativeFunctionLinearMapType(
             std::make_pair(paramType, i));
       }
       if (diffParam.isInOut()) {
+        if (paramType->isVoid())
+          continue;
         inoutParams.push_back(diffParam);
         continue;
       }

lib/SIL/IR/SILDeclRef.cpp

@@ -781,7 +781,13 @@ 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 originalFn =
+        getDecl()->getInterfaceType()->castTo<AnyFunctionType>();
+    SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+    autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+    auto numResults = semanticResults.size();
+    auto *resultIndices = IndexSubset::getDefault(
+        getDecl()->getASTContext(), numResults, /*includeAll*/ true);
     AutoDiffConfig silConfig(
         silParameterIndices, resultIndices,
         derivativeFunctionIdentifier->getDerivativeGenericSignature());

lib/Sema/TypeCheckAttr.cpp

@@ -4551,7 +4551,12 @@ IndexSubset *DifferentiableAttributeTypeCheckRequest::evaluate(
     }
     getterDecl->getAttrs().add(newAttr);
     // Register derivative function configuration.
-    auto *resultIndices = IndexSubset::get(ctx, 1, {0});
+    auto originalFn = getterDecl->getInterfaceType()->castTo<AnyFunctionType>();
+    SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+    autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+    auto numResults = semanticResults.size();
+    auto *resultIndices = IndexSubset::getDefault(
+        ctx, numResults, /*includeAll*/ true);
     getterDecl->addDerivativeFunctionConfiguration(
         {resolvedDiffParamIndices, resultIndices, derivativeGenSig});
     return resolvedDiffParamIndices;
@@ -4566,7 +4571,11 @@ IndexSubset *DifferentiableAttributeTypeCheckRequest::evaluate(
     return nullptr;
   }
   // Register derivative function configuration.
-  auto *resultIndices = IndexSubset::get(ctx, 1, {0});
+  SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+  autodiff::getFunctionSemanticResultTypes(originalFnRemappedTy, semanticResults);
+  auto numResults = semanticResults.size();
+  auto *resultIndices = IndexSubset::getDefault(
+      ctx, numResults, /*includeAll*/ true);
   original->addDerivativeFunctionConfiguration(
       {resolvedDiffParamIndices, resultIndices, derivativeGenSig});
   return resolvedDiffParamIndices;
@@ -5005,7 +5014,12 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
   }
 
   // Register derivative function configuration.
-  auto *resultIndices = IndexSubset::get(Ctx, 1, {0});
+  auto originalFn = originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
+  SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+  autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+  auto numResults = semanticResults.size();
+  auto *resultIndices = IndexSubset::getDefault(
+      Ctx, numResults, /*includeAll*/ true);
   originalAFD->addDerivativeFunctionConfiguration(
       {resolvedDiffParamIndices, resultIndices,
        derivative->getGenericSignature()});

lib/Sema/TypeCheckProtocol.cpp

@@ -457,7 +457,13 @@ matchWitnessDifferentiableAttr(DeclContext *dc, ValueDecl *req,
           witness->getAttrs().add(newAttr);
           success = true;
           // Register derivative function configuration.
-          auto *resultIndices = IndexSubset::get(ctx, 1, {0});
+          auto originalFn =
+              witnessAFD->getInterfaceType()->castTo<AnyFunctionType>();
+          SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+          autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+          auto numResults = semanticResults.size();
+          auto *resultIndices = IndexSubset::getDefault(
+              ctx, numResults, /*includeAll*/ true);
           witnessAFD->addDerivativeFunctionConfiguration(
               {newAttr->getParameterIndices(), resultIndices,
                newAttr->getDerivativeGenericSignature()});

lib/Serialization/ModuleFile.cpp

@@ -654,9 +654,14 @@ void ModuleFile::loadDerivativeFunctionConfigurations(
     }
     auto derivativeGenSig = derivativeGenSigOrError.get();
     // NOTE(TF-1038): Result indices are currently unsupported in derivative
-    // registration attributes. In the meantime, always use `{0}` (wrt the
-    // first and only result).
-    auto resultIndices = IndexSubset::get(ctx, 1, {0});
+    // registration attributes. In the meantime, always use all results.
+    auto originalFn =
+        originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
+    SmallVector<AutoDiffSemanticFunctionResultType, 1> semanticResults;
+    autodiff::getFunctionSemanticResultTypes(originalFn, semanticResults);
+    auto numResults = semanticResults.size();
+    auto *resultIndices = IndexSubset::getDefault(
+        ctx, numResults, /*includeAll*/ true);
     results.insert({parameterIndices, resultIndices, derivativeGenSig});
   }
 }