[AutoDiff upstream] Add TangentSpace. (#29107)

`TangentSpace` represents the tangent space of a type.

- For `Differentiable`-conforming types:
  - The tangent space is the `TangentVector` associated type.
- For tuple types:
  - The tangent space is a tuple of the elements' tangent space types, for the
    elements that have a tangent space.
- Other types have no tangent space.

`TypeBase::getAutoDiffTangentSpace` gets the tangent space of a type.

`TangentSpace` is used to:
- Compute the derivative function type of a given original function type.
- Compute the type of tangent/adjoint values during automatic differentiation.

Progress towards TF-828: upstream `@differentiable` attribute type-checking.

Author: dan-zheng

include/swift/AST/ASTContext.h

@@ -284,6 +284,9 @@ class ASTContext final {
   /// across invocations of both the parser and the type-checker.
   unsigned NextAutoClosureDiscriminator = 0;
 
+  /// Cached mapping from types to their associated tangent spaces.
+  llvm::DenseMap<Type, Optional<TangentSpace>> AutoDiffTangentSpaces;
+
   /// Cache of `@derivative` attributes keyed by parameter indices and
   /// derivative function kind. Used to diagnose duplicate `@derivative`
   /// attributes for the same key.

include/swift/AST/AutoDiff.h

@@ -29,6 +29,7 @@
 namespace swift {
 
 class AnyFunctionType;
+class TupleType;
 
 /// A function type differentiability kind.
 enum class DifferentiabilityKind : uint8_t {
@@ -153,6 +154,74 @@ class ParsedAutoDiffParameter {
   }
 };
 
+/// The tangent space of a type.
+///
+/// For `Differentiable`-conforming types:
+/// - The tangent space is the `TangentVector` associated type.
+///
+/// For tuple types:
+/// - The tangent space is a tuple of the elements' tangent space types, for the
+///   elements that have a tangent space.
+///
+/// Other types have no tangent space.
+class TangentSpace {
+public:
+  /// A tangent space kind.
+  enum class Kind {
+    /// The `TangentVector` associated type of a `Differentiable`-conforming
+    /// type.
+    TangentVector,
+    /// A product of tangent spaces as a tuple.
+    Tuple
+  };
+
+private:
+  Kind kind;
+  union Value {
+    // TangentVector
+    Type tangentVectorType;
+    // Tuple
+    TupleType *tupleType;
+
+    Value(Type tangentVectorType) : tangentVectorType(tangentVectorType) {}
+    Value(TupleType *tupleType) : tupleType(tupleType) {}
+  } value;
+
+  TangentSpace(Kind kind, Value value) : kind(kind), value(value) {}
+
+public:
+  TangentSpace() = delete;
+
+  static TangentSpace getTangentVector(Type tangentVectorType) {
+    return {Kind::TangentVector, tangentVectorType};
+  }
+  static TangentSpace getTuple(TupleType *tupleTy) {
+    return {Kind::Tuple, tupleTy};
+  }
+
+  bool isTangentVector() const { return kind == Kind::TangentVector; }
+  bool isTuple() const { return kind == Kind::Tuple; }
+
+  Kind getKind() const { return kind; }
+  Type getTangentVector() const {
+    assert(kind == Kind::TangentVector);
+    return value.tangentVectorType;
+  }
+  TupleType *getTuple() const {
+    assert(kind == Kind::Tuple);
+    return value.tupleType;
+  }
+
+  /// Get the tangent space type.
+  Type getType() const;
+
+  /// Get the tangent space canonical type.
+  CanType getCanonicalType() const;
+
+  /// Get the underlying nominal type declaration of the tangent space type.
+  NominalTypeDecl *getNominal() const;
+};
+
 /// Automatic differentiation utility namespace.
 namespace autodiff {
 

include/swift/AST/Types.h

@@ -1161,6 +1161,11 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   /// object type.
   TypeTraitResult canBeClass();
 
+  /// Return the tangent space of the given type, if it exists. Otherwise,
+  /// return `None`.
+  Optional<TangentSpace>
+  getAutoDiffTangentSpace(LookupConformanceFn lookupConformance);
+
 private:
   // Make vanilla new/delete illegal for Types.
   void *operator new(size_t Bytes) throw() = delete;

lib/AST/AutoDiff.cpp

@@ -66,3 +66,21 @@ void autodiff::getSubsetParameterTypes(IndexSubset *subset,
         results.push_back(curryLevel->getParams()[paramIndex].getOldType());
   }
 }
+
+Type TangentSpace::getType() const {
+  switch (kind) {
+  case Kind::TangentVector:
+    return value.tangentVectorType;
+  case Kind::Tuple:
+    return value.tupleType;
+  }
+}
+
+CanType TangentSpace::getCanonicalType() const {
+  return getType()->getCanonicalType();
+}
+
+NominalTypeDecl *TangentSpace::getNominal() const {
+  assert(isTangentVector());
+  return getTangentVector()->getNominalOrBoundGenericNominal();
+}

lib/AST/Type.cpp

@@ -4855,6 +4855,61 @@ CanType swift::substOpaqueTypesWithUnderlyingTypes(CanType ty,
   return ty.subst(replacer, replacer, flags)->getCanonicalType();
 }
 
+Optional<TangentSpace>
+TypeBase::getAutoDiffTangentSpace(LookupConformanceFn lookupConformance) {
+  assert(lookupConformance);
+  auto &ctx = getASTContext();
+
+  Type cacheKey = this;
+  auto lookup = ctx.AutoDiffTangentSpaces.find(cacheKey);
+  if (lookup != ctx.AutoDiffTangentSpaces.end())
+    return lookup->getSecond();
+  auto cache = [&](Optional<TangentSpace> tangentSpace) {
+    ctx.AutoDiffTangentSpaces.insert({cacheKey, tangentSpace});
+    return tangentSpace;
+  };
+
+  // For tuple types: the tangent space is a tuple of the elements'  tangent
+  // space types, for the elements that have a tangent space.
+  if (auto *tupleTy = getAs<TupleType>()) {
+    SmallVector<TupleTypeElt, 8> newElts;
+    for (auto elt : tupleTy->getElements()) {
+      auto eltSpace = elt.getType()->getAutoDiffTangentSpace(lookupConformance);
+      if (!eltSpace)
+        continue;
+      newElts.push_back(elt.getWithType(eltSpace->getType()));
+    }
+    if (newElts.empty())
+      return cache(
+          TangentSpace::getTuple(ctx.TheEmptyTupleType->castTo<TupleType>()));
+    if (newElts.size() == 1)
+      return cache(TangentSpace::getTangentVector(newElts.front().getType()));
+    auto *tupleType = TupleType::get(newElts, ctx)->castTo<TupleType>();
+    return cache(TangentSpace::getTuple(tupleType));
+  }
+
+  // For `Differentiable`-conforming types: the tangent space is the
+  // `TangentVector` associated type.
+  auto *differentiableProtocol =
+      ctx.getProtocol(KnownProtocolKind::Differentiable);
+  assert(differentiableProtocol && "`Differentiable` protocol not found");
+  auto associatedTypeLookup =
+      differentiableProtocol->lookupDirect(ctx.Id_TangentVector);
+  assert(associatedTypeLookup.size() == 1);
+  auto *dependentType = DependentMemberType::get(
+      differentiableProtocol->getDeclaredInterfaceType(),
+      cast<AssociatedTypeDecl>(associatedTypeLookup[0]));
+
+  // Try to get the `TangentVector` associated type of `base`.
+  // Return the associated type if it is valid.
+  auto assocTy = dependentType->substBaseType(this, lookupConformance);
+  if (!assocTy->hasError())
+    return cache(TangentSpace::getTangentVector(assocTy));
+
+  // Otherwise, there is no associated tangent space. Return `None`.
+  return cache(None);
+}
+
 CanSILFunctionType
 SILFunctionType::withSubstitutions(SubstitutionMap subs) const {
   return SILFunctionType::get(getSubstGenericSignature(),