[AutoDiff] WIP - Adding support for differentiating wrt inout parameters.

include/swift/AST/DiagnosticsSema.def

@@ -2738,7 +2738,11 @@ ERROR(implements_attr_protocol_not_conformed_to,none,
 // SWIFT_ENABLE_TENSORFLOW
 // @differentiable
 ERROR(differentiable_attr_void_result,none,
-      "cannot differentiate void function %0", (DeclName))
+      "cannot differentiate void function %0 with no inout wrt parameters", (DeclName))
+ERROR(differentiable_attr_non_void_inout,none,
+      "cannot differentiate non-void function %0 with inout wrt parameter", (DeclName))
+ERROR(differentiable_attr_multiple_inout,none,
+      "cannot differentiate function %0 with multiple inout wrt parameters", (DeclName))
 ERROR(differentiable_attr_associated_function_protocol,none,
       "cannot specify associated differentiation function on protocol "
       "requirement", ())

lib/AST/Type.cpp

@@ -4499,7 +4499,7 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
   SmallVector<AnyFunctionType *, 2> curryLevels;
   auto *currentLevel = eraseDynamicSelfType()->castTo<AnyFunctionType>();
   for (unsigned i : range(2)) {
-    (void)i;
+    (void) i;
     if (currentLevel == nullptr)
       break;
     curryLevels.push_back(currentLevel);
@@ -4508,6 +4508,38 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
 
   Type originalResult = curryLevels.back()->getResult();
 
+  // Perform validity checks for `inout` parameters. If the function is being
+  // differentiated with respect to more than one `inout` parameters, we
+  // produce an error. Also, if the original function type has a `Void` return
+  // type, we try to see if any `inout` parameter is being differentiated with
+  // respect to. If not, we produce an error. For example:
+  //   - (T0, T1, T2) -> Void, wrt: 0, 1              // bad: no `inout` wrt parameters.
+  //   - (inout T0, inout T1, T2) -> Void, wrt: 0, 1  // bad: more than one `inout` wrt parameters.
+  //   - (T0, inout T1, T2) -> Void, wrt: 0, 1        // good: exactly one `inout` wrt parameter.
+  Type* inoutWrtParamType = nullptr;
+  for (auto wrtParamType : wrtParamTypes) {
+    if (wrtParamType->is<InOutType>()) {
+      assert(inoutWrtParamType == nullptr &&
+             "no more than one `inout` wrt parameters are allowed");
+      inoutWrtParamType = &wrtParamType->getInOutObjectType();
+    }
+  }
+  assert((!originalResult->isVoid() || inoutWrtParamType != nullptr) &&
+         "no `inout` wrt parameters for function with `Void` return type");
+
+  // We also check if the original function has a return type, along with
+  // differentiating with respect to at least one `inout` parameter. For example:
+  //   - (T0, T1, T2) -> R       // good
+  //   - (T0, inout T1, T2) -> R // bad
+  assert((originalResult->isVoid() || inoutWrtParamType == nullptr) &&
+         "non-`Void` return type and differentiating wrt `inout` parameter");
+
+  // For inout arguments, we treat the inout type as both an argument type and
+  // a result type.
+  if (inoutWrtParamType != nullptr) {
+    originalResult = inoutWrtParamType;
+  }
+
   // Build the closure type, which is different depending on whether this is a
   // JVP or VJP.
   Type closure;
@@ -4518,20 +4550,21 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
     SmallVector<AnyFunctionType::Param, 8> differentialParams;
     for (auto wrtParamType : wrtParamTypes)
       differentialParams.push_back(
-          AnyFunctionType::Param(
-              wrtParamType->getAutoDiffAssociatedTangentSpace(lookupConformance)
-                  ->getType()));
+          AnyFunctionType::Param(wrtParamType->getInOutObjectType()
+               ->getAutoDiffAssociatedTangentSpace(lookupConformance)
+               ->getType()));
 
     SmallVector<TupleTypeElt, 8> differentialResults;
     if (auto *resultTuple = originalResult->getAs<TupleType>()) {
       auto resultTupleEltType = resultTuple->getElementType(resultIndex);
-      differentialResults.push_back(resultTupleEltType
-          ->getAutoDiffAssociatedTangentSpace(lookupConformance)->getType());
+      differentialResults.push_back(resultTupleEltType->getInOutObjectType()
+          ->getAutoDiffAssociatedTangentSpace(lookupConformance)
+          ->getType());
     } else {
       assert(resultIndex == 0 && "resultIndex out of bounds");
-      differentialResults.push_back(
-          originalResult->getAutoDiffAssociatedTangentSpace(lookupConformance)
-              ->getType());
+      differentialResults.push_back(originalResult->getInOutObjectType()
+          ->getAutoDiffAssociatedTangentSpace(lookupConformance)
+          ->getType());
     }
     Type differentialResult =
         differentialResults.size() > 1
@@ -4548,22 +4581,22 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
     if (auto *resultTuple = originalResult->getAs<TupleType>()) {
       auto resultTupleEltType = resultTuple->getElementType(resultIndex);
       pullbackParams.push_back(
-          AnyFunctionType::Param(resultTupleEltType
-              ->getAutoDiffAssociatedTangentSpace(lookupConformance)
-                  ->getType()));
+          AnyFunctionType::Param(resultTupleEltType->getInOutObjectType()
+               ->getAutoDiffAssociatedTangentSpace(lookupConformance)
+               ->getType()));
     } else {
       assert(resultIndex == 0 && "resultIndex out of bounds");
       pullbackParams.push_back(
-          AnyFunctionType::Param(originalResult
-              ->getAutoDiffAssociatedTangentSpace(lookupConformance)
-                  ->getType()));
+          AnyFunctionType::Param(originalResult->getInOutObjectType()
+               ->getAutoDiffAssociatedTangentSpace(lookupConformance)
+               ->getType()));
     }
 
     SmallVector<TupleTypeElt, 8> pullbackResults;
     for (auto wrtParamType : wrtParamTypes)
-      pullbackResults.push_back(wrtParamType
+      pullbackResults.push_back(wrtParamType->getInOutObjectType()
           ->getAutoDiffAssociatedTangentSpace(lookupConformance)
-              ->getType());
+          ->getType());
     Type pullbackResult = pullbackResults.size() > 1
                               ? TupleType::get(pullbackResults, ctx)
                               : pullbackResults[0].getType();

lib/SIL/SILType.cpp

@@ -595,6 +595,13 @@ bool SILType::isLoweringOf(SILModule &Mod, CanType formalType) {
 // SWIFT_ENABLE_TENSORFLOW
 /// Returns true if this SILType is a differentiable type.
 bool SILType::isDifferentiable(SILModule &M) const {
-  return getASTType()->getAutoDiffAssociatedTangentSpace(
-      LookUpConformanceInModule(M.getSwiftModule())).hasValue();
+  auto lookupConformance = LookUpConformanceInModule(M.getSwiftModule());
+  auto ty = getASTType();
+  if (!ty->is<InOutType>()) {
+    return ty->getAutoDiffAssociatedTangentSpace(
+        lookupConformance).hasValue();
+  }
+  Type base = ty->getInOutObjectType();
+  return base->getAutoDiffAssociatedTangentSpace(
+      lookupConformance).hasValue();
 }

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -3136,11 +3136,12 @@ class VJPEmitter final
     // Check for active 'inout' arguments.
     auto paramInfos = ai->getSubstCalleeConv().getParameters();
     bool hasActiveInoutParams = false;
-    for (unsigned i : swift::indices(paramInfos))
+    for (unsigned i : swift::indices(paramInfos)) {
       if (paramInfos[i].isIndirectInOut() &&
           activityInfo.isActive(ai->getArgumentsWithoutIndirectResults()[i],
                                 getIndices()))
         hasActiveInoutParams = true;
+    }
     // Reject functions with active inout arguments. It's not yet supported.
     if (hasActiveInoutParams) {
       context.emitNondifferentiabilityError(ai, invoker,