[AutoDiff] Update TypeBase::getAutoDiffAssociatedVectorSpace to handle function types. (#22166)

Formally, when a type `T where T : Differentiable` gets abstracted as a function `(X...) -> T` for any `X...`, the differentiability of the abstracted type depends entirely on the differentiability of `T`. Since structural types cannot conform to protocols yet in Swift, we need to handle in AD-associated type calculation the same way we handle tuples.

The type calculation rules are better described as code, in imaginary syntax where parameterized extensions, variadic generic parameters, and protocol conformances for structural types are supported.

```swift
extension<T..., U> ((T...) -> U) : Differentiable where U : Differentiable {
  public typealias TangentVector = (T...) -> U.TangentVector
  public typealias CotangentVector = (T...) -> U.CotangentVector
  public func moved(along direction: TangentVector) -> (T...) -> U {
    return { (x...) in self(x...).moved(along: direction(x...)) }
  }
  public func tangentVector(from cotangent: CotangentVector) -> TangentVector {
    return { (x...) in self(x...).tangentVector(from: cotangent(x...)) }
  }
}
```

This is a crucial step towards the correct typing of curried differentiable functions, which helps us differentiate through curry thunks for methods.

```swift
func curry<T : Differentiable, U : Differentiable>(
  f: @autodiff (T, U) -> V
) -> @autodiff (T) -> @autodiff (U) -> V {
  return { x in { y in f(x, y) } }
}
```

Partially resolves [SR-9448](https://bugs.swift.org/browse/SR-9448), which needs this patch to be able to calculate the associate vector space of a curried function.

Author: rxwei

include/swift/AST/AutoDiff.h

@@ -359,10 +359,12 @@ class VectorSpace {
 public:
   /// A tangent space kind.
   enum class Kind {
-    /// A type that conforms to `VectorNumeric`.
+    /// A type that conforms to `AdditiveArithmetic`.
     Vector,
     /// A product of vector spaces as a tuple.
     Tuple,
+    /// A function type whose innermost result conforms to `AdditiveArithmetic`.
+    Function
   };
 
 private:
@@ -372,9 +374,12 @@ class VectorSpace {
     Type vectorType;
     // Tuple
     TupleType *tupleType;
+    // Function
+    AnyFunctionType *functionType;
 
     Value(Type vectorType) : vectorType(vectorType) {}
     Value(TupleType *tupleType) : tupleType(tupleType) {}
+    Value(AnyFunctionType *functionType) : functionType(functionType) {}
   } value;
 
   VectorSpace(Kind kind, Value value)
@@ -389,6 +394,9 @@ class VectorSpace {
   static VectorSpace getTuple(TupleType *tupleTy) {
     return {Kind::Tuple, tupleTy};
   }
+  static VectorSpace getFunction(AnyFunctionType *fnTy) {
+    return {Kind::Function, fnTy};
+  }
 
   bool isVector() const { return kind == Kind::Vector; }
   bool isTuple() const { return kind == Kind::Tuple; }
@@ -402,6 +410,10 @@ class VectorSpace {
     assert(kind == Kind::Tuple);
     return value.tupleType;
   }
+  AnyFunctionType *getFunction() const {
+    assert(kind == Kind::Function);
+    return value.functionType;
+  }
 
   Type getType() const;
   CanType getCanonicalType() const;

lib/AST/AutoDiff.cpp

@@ -289,6 +289,8 @@ Type VectorSpace::getType() const {
     return value.vectorType;
   case Kind::Tuple:
     return value.tupleType;
+  case Kind::Function:
+    return value.functionType;
   }
 }
 

lib/AST/Type.cpp

@@ -4089,6 +4089,21 @@ Type TypeBase::openAnyExistentialType(ArchetypeType *&opened) {
 }
 
 // SWIFT_ENABLE_TENSORFLOW
+// Makes a function with the same generic signature and extinfo as `copy`, but
+// with `params` parameters and `retTy` return type.
+static AnyFunctionType *
+makeFunctionType(AnyFunctionType *copy, ArrayRef<AnyFunctionType::Param> params,
+                 Type retTy, GenericSignature *whereClauseGenSig) {
+  auto genericSignature = whereClauseGenSig;
+  if (!genericSignature)
+    if (auto *genericFunctionType = copy->getAs<GenericFunctionType>())
+      genericSignature = genericFunctionType->getGenericSignature();
+  if (genericSignature)
+    return GenericFunctionType::get(genericSignature, params, retTy,
+                                    copy->getExtInfo());
+  return FunctionType::get(params, retTy, copy->getExtInfo());
+}
+
 Optional<VectorSpace> TypeBase::getAutoDiffAssociatedVectorSpace(
     AutoDiffAssociatedVectorSpaceKind kind,
     LookupConformanceFn lookupConformance) {
@@ -4104,7 +4119,19 @@ Optional<VectorSpace> TypeBase::getAutoDiffAssociatedVectorSpace(
     return vs;
   };
 
-  // Tuples' Tangent/Cotangent is a tuple of each element's Tangent/Cotangent.
+  // Functions' tangent/cotangent is the same function except the innermost
+  // return type being replaced by its tangent/cotangent.
+  if (auto *fnTy = getAs<AnyFunctionType>()) {
+    auto resultSpace = fnTy->getResult()->getAutoDiffAssociatedVectorSpace(
+        kind, lookupConformance);
+    if (!resultSpace)
+      return cache(None);
+    return VectorSpace::getFunction(
+        makeFunctionType(fnTy, fnTy->getParams(), resultSpace->getType(),
+                         fnTy->getOptGenericSignature()));
+  }
+
+  // Tuples' tangent/cotangent is a tuple of each element's Tangent/Cotangent.
   if (auto *tupleTy = getAs<TupleType>()) {
     SmallVector<TupleTypeElt, 8> newElts;
     for (auto elt : tupleTy->getElements()) {
@@ -4148,21 +4175,6 @@ Optional<VectorSpace> TypeBase::getAutoDiffAssociatedVectorSpace(
   return cache(None);
 }
 
-// Makes a function with the same generic signature and extinfo as `copy`, but
-// with `params` parameters and `retTy` return type.
-static AnyFunctionType *
-makeFunctionType(AnyFunctionType *copy, ArrayRef<AnyFunctionType::Param> params,
-                 Type retTy, GenericSignature *whereClauseGenSig) {
-  auto genericSignature = whereClauseGenSig;
-  if (!genericSignature)
-    if (auto *genericFunctionType = copy->getAs<GenericFunctionType>())
-      genericSignature = genericFunctionType->getGenericSignature();
-  if (genericSignature)
-    return GenericFunctionType::get(genericSignature, params, retTy,
-                                    copy->getExtInfo());
-  return FunctionType::get(params, retTy, copy->getExtInfo());
-}
-
 AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
     AutoDiffParameterIndices *indices, unsigned resultIndex,
     unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -3799,6 +3799,10 @@ SILValue AdjointEmitter::accumulateMaterializedAdjointsDirect(SILValue lhs,
     }
     return builder.createTuple(loc, adjointTy, adjElements);
   }
+  case VectorSpace::Kind::Function: {
+    llvm_unreachable(
+      "Unimplemented: Emit thunks for abstracting adjoint accumulation");
+  }
   }
 }
 
@@ -3854,6 +3858,10 @@ void AdjointEmitter::accumulateMaterializedAdjointsIndirect(
     }
     return;
   }
+  case VectorSpace::Kind::Function: {
+    llvm_unreachable(
+        "Unimplemented: Emit thunks for abstracting adjoint accumulation");
+  }
   }
 }