[AutoDiff] Support differentiation of apply with inout arguments.

Add reverse-mode differentiation support for `apply` with `inout` arguments.

Notable pullback generation changes:
- If the pullback seed argument is `inout`, assign it (rather than a copy)
  directly as the adjoint buffer of the original result. This is important so
  the value is updated in-place.
- In `visitApplyInst`: skip adjoint accumulation for `inout` arguments.
  Adjoint accumulation for `inout` arguments occurs when callee pullbacks are
  applied, so no extra accumulation is necessary.

Add derivatives for functions with `inout` parameters in the stdlib for testing:
- `FloatingPoint` operations: `+=`, `-=`, `*=`, `/=`
- `Array.append`

Resolves TF-1165.

Todos:
- Add more tests, e.g. SILGen tests for `inout` derivative typing rules.
- Evaluate performance of `inout` derivatives vs functional derivatives + mutation.
- TF-1166: enable `@differentiable` attribute on `set` accessors.
- TF-1173: add forward-mode differentiation support for `apply` with `inout`
  parameters.

Author: dan-zheng

include/swift/SIL/AbstractionPattern.h

@@ -1146,8 +1146,7 @@ class AbstractionPattern {
   /// The arguments are the same as the arguments to
   /// `AnyFunctionType::getAutoDiffDerivativeFunctionType()`.
   AbstractionPattern getAutoDiffDerivativeFunctionType(
-      IndexSubset *indices, unsigned resultIndex,
-      AutoDiffDerivativeFunctionKind kind,
+      IndexSubset *indices, AutoDiffDerivativeFunctionKind kind,
       LookupConformanceFn lookupConformance,
       GenericSignature derivativeGenericSignature = GenericSignature(),
       bool makeSelfParamFirst = false);

lib/AST/Builtins.cpp

@@ -1035,7 +1035,7 @@ static ValueDecl *getAutoDiffApplyDerivativeFunction(
   BuiltinFunctionBuilder::LambdaGenerator resultGen{
       [=, &Context](BuiltinFunctionBuilder &builder) -> Type {
         auto derivativeFnTy = diffFnType->getAutoDiffDerivativeFunctionType(
-            paramIndices, /*resultIndex*/ 0, kind,
+            paramIndices, kind,
             LookUpConformanceInModule(Context.TheBuiltinModule));
         return derivativeFnTy->getResult();
       }};

lib/SIL/AbstractionPattern.cpp

@@ -926,17 +926,17 @@ const {
 }
 
 AbstractionPattern AbstractionPattern::getAutoDiffDerivativeFunctionType(
-    IndexSubset *indices, unsigned resultIndex,
-    AutoDiffDerivativeFunctionKind kind, LookupConformanceFn lookupConformance,
+    IndexSubset *indices, AutoDiffDerivativeFunctionKind kind,
+    LookupConformanceFn lookupConformance,
     GenericSignature derivativeGenericSignature, bool makeSelfParamFirst) {
   switch (getKind()) {
   case Kind::Type: {
     auto fnTy = dyn_cast<AnyFunctionType>(getType());
     if (!fnTy)
       return getOpaqueDerivativeFunction();
     auto derivativeFnTy = fnTy->getAutoDiffDerivativeFunctionType(
-        indices, resultIndex, kind, lookupConformance,
-        derivativeGenericSignature, makeSelfParamFirst);
+        indices, kind, lookupConformance, derivativeGenericSignature,
+        makeSelfParamFirst);
     assert(derivativeFnTy);
     return AbstractionPattern(getGenericSignature(),
                               derivativeFnTy->getCanonicalType());

lib/SIL/SILFunctionType.cpp

@@ -231,7 +231,8 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
     LookupConformanceFn lookupConformance,
     CanGenericSignature derivativeFnGenSig, bool isReabstractionThunk) {
   auto &ctx = getASTContext();
-  auto resultIndices = IndexSubset::get(ctx, getNumResults(), {resultIndex});
+  auto *resultIndices = IndexSubset::get(
+      ctx, getNumResults() + getNumIndirectMutatingParameters(), {resultIndex});
   SILAutoDiffDerivativeFunctionKey key{
       this, parameterIndices,   resultIndices,
       kind, derivativeFnGenSig, isReabstractionThunk};

lib/SIL/TypeLowering.cpp

@@ -2181,8 +2181,8 @@ CanAnyFunctionType TypeConverter::makeConstantInterfaceType(SILDeclRef c) {
     auto originalFnTy =
         makeConstantInterfaceType(c.asAutoDiffOriginalFunction());
     auto *fnTy = originalFnTy->getAutoDiffDerivativeFunctionType(
-        autoDiffFuncId->getParameterIndices(), /*resultIndex*/ 0,
-        autoDiffFuncId->getKind(), LookUpConformanceInModule(&M));
+        autoDiffFuncId->getParameterIndices(), autoDiffFuncId->getKind(),
+        LookUpConformanceInModule(&M));
     return cast<AnyFunctionType>(fnTy->getCanonicalType());
   }
 

lib/SILGen/SILGenPoly.cpp

@@ -3342,18 +3342,18 @@ static ManagedValue createAutoDiffThunk(SILGenFunction &SGF,
   auto *parameterIndices = IndexSubset::get(SGF.getASTContext(), parameterBits);
 
   auto getDerivativeFnTy =
-      [&](CanAnyFunctionType fnTy, AutoDiffDerivativeFunctionKind kind)
-          -> CanAnyFunctionType {
-        auto assocTy = fnTy->getAutoDiffDerivativeFunctionType(
-            parameterIndices, /*resultIndex*/ 0,
-            kind, LookUpConformanceInModule(SGF.SGM.M.getSwiftModule()));
-        return cast<AnyFunctionType>(assocTy->getCanonicalType());
-      };
+      [&](CanAnyFunctionType fnTy,
+          AutoDiffDerivativeFunctionKind kind) -> CanAnyFunctionType {
+    auto assocTy = fnTy->getAutoDiffDerivativeFunctionType(
+        parameterIndices, kind,
+        LookUpConformanceInModule(SGF.SGM.M.getSwiftModule()));
+    return cast<AnyFunctionType>(assocTy->getCanonicalType());
+  };
   auto getDerivativeFnPattern =
       [&](AbstractionPattern pattern,
           AutoDiffDerivativeFunctionKind kind) -> AbstractionPattern {
     return pattern.getAutoDiffDerivativeFunctionType(
-        parameterIndices, /*resultIndex*/ 0, kind,
+        parameterIndices, kind,
         LookUpConformanceInModule(SGF.SGM.M.getSwiftModule()));
   };
   auto createDerivativeFnThunk =

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -508,9 +508,10 @@ emitDerivativeFunctionReference(
     auto *originalFn = originalFRI->getReferencedFunctionOrNull();
     assert(originalFn);
     auto originalFnTy = originalFn->getLoweredFunctionType();
-    auto *desiredResultIndices =
-        IndexSubset::get(context.getASTContext(), originalFnTy->getNumResults(),
-                         {desiredIndices.source});
+    auto numResults = originalFnTy->getNumResults() +
+                      originalFnTy->getNumIndirectMutatingParameters();
+    auto *desiredResultIndices = IndexSubset::get(
+        context.getASTContext(), numResults, {desiredIndices.source});
     auto *desiredParameterIndices = desiredIndices.parameters;
     // NOTE(TF-893): Extending capacity is necessary when `originalFnTy` has
     // parameters corresponding to captured variables.
@@ -550,9 +551,19 @@ emitDerivativeFunctionReference(
         }
       }
       // Check and diagnose non-differentiable results.
-      if (!originalFnTy->getResults()[desiredIndices.source]
-               .getSILStorageInterfaceType()
-               .isDifferentiable(context.getModule())) {
+      SILType resultType;
+      if (desiredIndices.source >= originalFnTy->getNumResults()) {
+        auto inoutParamIdx =
+            desiredIndices.source - originalFnTy->getNumResults();
+        auto inoutParam =
+            *std::next(originalFnTy->getIndirectMutatingParameters().begin(),
+                       inoutParamIdx);
+        resultType = inoutParam.getSILStorageInterfaceType();
+      } else {
+        resultType = originalFnTy->getResults()[desiredIndices.source]
+                         .getSILStorageInterfaceType();
+      }
+      if (!resultType.isDifferentiable(context.getModule())) {
         context.emitNondifferentiabilityError(
             original, invoker, diag::autodiff_nondifferentiable_result);
         return None;

lib/SILOptimizer/Utils/Differentiation/Common.cpp

@@ -119,6 +119,15 @@ void collectAllFormalResultsInTypeOrder(SILFunction &function,
   for (auto &resInfo : convs.getResults())
     results.push_back(resInfo.isFormalDirect() ? dirResults[dirResIdx++]
                                                : indResults[indResIdx++]);
+  // Treat `inout` parameters as semantic results.
+  // Append `inout` parameters after formal results.
+  for (auto i : range(convs.getNumParameters())) {
+    auto paramInfo = convs.getParameters()[i];
+    if (!paramInfo.isIndirectMutating())
+      continue;
+    auto *argument = function.getArgumentsWithoutIndirectResults()[i];
+    results.push_back(argument);
+  }
 }
 
 /// Given a function, gathers all of its direct results in an order defined by
@@ -194,8 +203,21 @@ void collectMinimalIndicesForFunctionCall(
       ++indResIdx;
     }
   }
+  // Record all `inout` parameters as results.
+  auto inoutParamResultIndex = calleeFnTy->getNumResults();
+  for (auto &paramAndIdx : enumerate(calleeConvs.getParameters())) {
+    auto &param = paramAndIdx.value();
+    if (!param.isIndirectMutating())
+      continue;
+    unsigned idx = paramAndIdx.index();
+    auto inoutArg = ai->getArgument(idx);
+    results.push_back(inoutArg);
+    resultIndices.push_back(inoutParamResultIndex++);
+  }
   // Make sure the function call has active results.
-  assert(results.size() == calleeFnTy->getNumResults());
+  auto numResults = calleeFnTy->getNumResults() +
+                    calleeFnTy->getNumIndirectMutatingParameters();
+  assert(results.size() == numResults);
   assert(llvm::any_of(results, [&](SILValue result) {
     return activityInfo.isActive(result, parentIndices);
   }));

lib/SILOptimizer/Utils/Differentiation/JVPEmitter.cpp

@@ -1000,10 +1000,12 @@ void JVPEmitter::prepareForDifferentialGeneration() {
   // Initialize tangent mapping for indirect results.
   auto origIndResults = original->getIndirectResults();
   auto diffIndResults = differential.getIndirectResults();
-  assert(origIndResults.size() == diffIndResults.size());
-
+  size_t numInoutParameters = llvm::count_if(
+      original->getLoweredFunctionType()->getParameters(),
+      [](SILParameterInfo paramInfo) { return paramInfo.isIndirectInOut(); });
+  assert(origIndResults.size() + numInoutParameters == diffIndResults.size());
   for (auto &origBB : *original)
-    for (auto i : indices(diffIndResults))
+    for (auto i : indices(origIndResults))
       setTangentBuffer(&origBB, origIndResults[i], diffIndResults[i]);
 }
 
@@ -1036,14 +1038,29 @@ JVPEmitter::createEmptyDifferential(ADContext &context,
   auto indices = witness->getSILAutoDiffIndices();
 
   // Add differential results.
-  auto origResult = origTy->getResults()[indices.source];
-  origResult = origResult.getWithInterfaceType(
-      origResult.getInterfaceType()->getCanonicalType(witnessCanGenSig));
-  dfResults.push_back(
-      SILResultInfo(origResult.getInterfaceType()
-                        ->getAutoDiffTangentSpace(lookupConformance)
-                        ->getCanonicalType(),
-                    origResult.getConvention()));
+  Optional<SILParameterInfo> inoutDiffParam = None;
+  for (auto origParam : origTy->getParameters()) {
+    if (!origParam.isIndirectInOut())
+      continue;
+    inoutDiffParam = origParam;
+  }
+
+  if (inoutDiffParam) {
+    dfResults.push_back(
+        SILResultInfo(inoutDiffParam->getInterfaceType()
+                          ->getAutoDiffTangentSpace(lookupConformance)
+                          ->getCanonicalType(),
+                      ResultConvention::Indirect));
+  } else {
+    auto origResult = origTy->getResults()[indices.source];
+    origResult = origResult.getWithInterfaceType(
+        origResult.getInterfaceType()->getCanonicalType(witnessCanGenSig));
+    dfResults.push_back(
+        SILResultInfo(origResult.getInterfaceType()
+                          ->getAutoDiffTangentSpace(lookupConformance)
+                          ->getCanonicalType(),
+                      origResult.getConvention()));
+  }
 
   // Add differential parameters for the requested wrt parameters.
   for (auto i : indices.parameters->getIndices()) {

lib/SILOptimizer/Utils/Differentiation/LinearMapInfo.cpp

@@ -244,8 +244,13 @@ VarDecl *LinearMapInfo::addLinearMapDecl(ApplyInst *ai, SILType linearMapType) {
   // the same parameters and results.
   auto silFnTy = linearMapType.castTo<SILFunctionType>();
   SmallVector<AnyFunctionType::Param, 8> params;
-  for (auto &param : silFnTy->getParameters())
-    params.push_back(AnyFunctionType::Param(param.getInterfaceType()));
+  for (auto &param : silFnTy->getParameters()) {
+    ParameterTypeFlags flags;
+    if (param.isIndirectMutating())
+      flags = flags.withInOut(true);
+    params.push_back(
+        AnyFunctionType::Param(param.getInterfaceType(), Identifier(), flags));
+  }
   AnyFunctionType *astFnTy;
   if (auto genSig = silFnTy->getSubstGenericSignature())
     astFnTy = GenericFunctionType::get(
@@ -283,11 +288,22 @@ void LinearMapInfo::addLinearMapToStruct(ADContext &context, ApplyInst *ai,
   auto hasActiveResults = llvm::any_of(allResults, [&](SILValue res) {
     return activityInfo.isActive(res, indices);
   });
-  auto hasActiveArguments =
-      llvm::any_of(ai->getArgumentsWithoutIndirectResults(), [&](SILValue arg) {
-        return activityInfo.isActive(arg, indices);
-      });
-  if (!hasActiveResults || !hasActiveArguments)
+  bool hasActiveInoutArgument = false;
+  bool hasActiveArguments = false;
+  auto numIndirectResults = ai->getNumIndirectResults();
+  for (auto argIdx : range(ai->getSubstCalleeConv().getNumParameters())) {
+    auto arg = ai->getArgumentsWithoutIndirectResults()[argIdx];
+    if (activityInfo.isActive(arg, indices)) {
+      hasActiveArguments = true;
+      auto paramInfo = ai->getSubstCalleeConv().getParamInfoForSILArg(
+          numIndirectResults + argIdx);
+      if (paramInfo.isIndirectMutating())
+        hasActiveInoutArgument = true;
+    }
+  }
+  if (!hasActiveArguments)
+    return;
+  if (!hasActiveResults && !hasActiveInoutArgument)
     return;
 
   // Compute differentiation result index.
@@ -323,8 +339,16 @@ void LinearMapInfo::addLinearMapToStruct(ADContext &context, ApplyInst *ai,
             return true;
         }
         // Check non-differentiable results.
-        auto remappedResultType = origFnTy->getResults()[applyIndices.source]
-                                      .getSILStorageInterfaceType();
+        SILType remappedResultType;
+        if (applyIndices.source >= origFnTy->getNumResults()) {
+          auto inoutArgIdx = applyIndices.source - origFnTy->getNumResults();
+          auto inoutArg =
+              *std::next(ai->getInoutArguments().begin(), inoutArgIdx);
+          remappedResultType = inoutArg->getType();
+        } else {
+          remappedResultType = origFnTy->getResults()[applyIndices.source]
+                                   .getSILStorageInterfaceType();
+        }
         if (!remappedResultType.isDifferentiable(derivative->getModule()))
           return true;
         return false;
@@ -338,13 +362,14 @@ void LinearMapInfo::addLinearMapToStruct(ADContext &context, ApplyInst *ai,
           parameters, source, derivativeFnKind, context.getTypeConverter(),
           LookUpConformanceInModule(derivative->getModule().getSwiftModule()));
 
-  auto derivativeFnResultTypes =
-      derivativeFnType->getAllResultsInterfaceType().castTo<TupleType>();
-  auto linearMapSILType = SILType::getPrimitiveObjectType(
-      derivativeFnResultTypes
-          ->getElement(derivativeFnResultTypes->getElements().size() - 1)
-          .getType()
-          ->getCanonicalType());
+  auto derivativeFnResultTypes = derivativeFnType->getAllResultsInterfaceType();
+  auto linearMapSILType = derivativeFnResultTypes;
+  if (auto tupleType = derivativeFnResultTypes.getAs<TupleType>()) {
+    linearMapSILType = SILType::getPrimitiveObjectType(
+        tupleType->getElement(tupleType->getElements().size() - 1)
+            .getType()
+            ->getCanonicalType());
+  }
   addLinearMapDecl(ai, linearMapSILType);
 }
 
@@ -396,21 +421,11 @@ void LinearMapInfo::generateDifferentiationDataStructures(
   for (auto &origBB : *original) {
     for (auto &inst : origBB) {
       if (auto *ai = dyn_cast<ApplyInst>(&inst)) {
-        // Skip `apply` instructions with active `inout` arguments.
-        // TODO(TF-129): Support `inout` argument differentiation.
-        bool hasActiveInoutArgument =
-            llvm::any_of(ai->getInoutArguments(), [&](SILValue inoutArg) {
-              return activityInfo.isActive(inoutArg, indices);
-            });
-        if (hasActiveInoutArgument)
-          continue;
-
         // Add linear map field to struct for active `apply` instructions.
         // Skip array literal intrinsic applications since array literal
         // initialization is linear and handled separately.
         if (!shouldDifferentiateApplySite(ai) || isArrayLiteralIntrinsic(ai))
           continue;
-
         LLVM_DEBUG(getADDebugStream()
                    << "Adding linear map struct field for " << *ai);
         addLinearMapToStruct(context, ai, indices);