[AutoDiff] NFC: Refactor TangentSpace (#21564)

* Fix SR-8770

* Refactor TangentSpace

Author: rxwei

include/swift/AST/ASTContext.h

@@ -104,7 +104,7 @@ namespace swift {
   class VarDecl;
   class UnifiedStatsReporter;
   // SWIFT_ENABLE_TENSORFLOW
-  class TangentSpace;
+  class VectorSpace;
 
   enum class KnownProtocolKind : uint8_t;
 
@@ -955,7 +955,7 @@ class ASTContext final {
   // SWIFT_ENABLE_TENSORFLOW
   /// Compute the tangent space of this manifold, if the given type represents a
   /// differentiable manifold.
-  Optional<TangentSpace> getTangentSpace(CanType type, ModuleDecl *module);
+  Optional<VectorSpace> getTangentSpace(CanType type, ModuleDecl *module);
 
 private:
   friend Decl;

include/swift/AST/AutoDiff.h

@@ -357,32 +357,29 @@ class StructDecl;
 class TupleType;
 class EnumDecl;
 
-/// A type that represents the tangent space of a differentiable type.
-class TangentSpace {
+/// A type that represents a vector space.
+class VectorSpace {
 public:
   /// A tangent space kind.
   enum class Kind {
     /// `Builtin.FP<...>`.
-    BuiltinRealScalar,
-    /// A type that conforms to `FloatingPoint`.
-    RealScalar,
-    /// A type that conforms to `VectorNumeric` where the associated
-    /// `ScalarElement` conforms to `FloatingPoint`.
-    RealVector,
-    /// A product of tangent spaces as a struct.
-    ProductStruct,
-    /// A product of tangent spaces as a tuple.
-    ProductTuple,
-    /// A sum of tangent spaces.
-    Sum
+    BuiltinFloat,
+    /// A type that conforms to `VectorNumeric`.
+    Vector,
+    /// A product of vector spaces as a struct.
+    Struct,
+    /// A product of vector spaces as a tuple.
+    Tuple,
+    /// A union of vector spaces.
+    Enum
   };
 
 private:
   Kind kind;
   union Value {
     // BuiltinRealScalar
     BuiltinFloatType *builtinFPType;
-    // RealScalar or RealVector
+    // RealVector
     NominalTypeDecl *realNominalType;
     // ProductStruct
     StructDecl *structDecl;
@@ -398,67 +395,55 @@ class TangentSpace {
     Value(EnumDecl *enumDecl) : enumDecl(enumDecl) {}
   } value;
 
-  TangentSpace(Kind kind, Value value)
+  VectorSpace(Kind kind, Value value)
       : kind(kind), value(value) {}
 
 public:
-  TangentSpace() = delete;
+  VectorSpace() = delete;
 
-  static TangentSpace
+  static VectorSpace
   getBuiltinRealScalarSpace(BuiltinFloatType *builtinFP) {
-    return {Kind::BuiltinRealScalar, builtinFP};
+    return {Kind::BuiltinFloat, builtinFP};
   }
-  static TangentSpace getRealScalarSpace(NominalTypeDecl *typeDecl) {
-    return {Kind::RealScalar, typeDecl};
+  static VectorSpace getRealVectorSpace(NominalTypeDecl *typeDecl) {
+    return {Kind::Vector, typeDecl};
   }
-  static TangentSpace getRealVectorSpace(NominalTypeDecl *typeDecl) {
-    return {Kind::RealVector, typeDecl};
+  static VectorSpace getStruct(StructDecl *structDecl) {
+    return {Kind::Struct, structDecl};
   }
-  static TangentSpace getProductStruct(StructDecl *structDecl) {
-    return {Kind::ProductStruct, structDecl};
+  static VectorSpace getTuple(TupleType *tupleTy) {
+    return {Kind::Tuple, tupleTy};
   }
-  static TangentSpace getProductTuple(TupleType *tupleTy) {
-    return {Kind::ProductTuple, tupleTy};
-  }
-  static TangentSpace getSum(EnumDecl *enumDecl) {
-    return {Kind::Sum, enumDecl};
+  static VectorSpace getEnum(EnumDecl *enumDecl) {
+    return {Kind::Enum, enumDecl};
   }
 
-  bool isBuiltinRealScalarSpace() const {
-    return kind == Kind::BuiltinRealScalar;
+  bool isBuiltinFloat() const {
+    return kind == Kind::BuiltinFloat;
   }
-  bool isRealScalarSpace() const { return kind == Kind::RealScalar; }
-  bool isRealVectorSpace() const { return kind == Kind::RealVector; }
-  bool isProductStruct() const { return kind == Kind::ProductStruct; }
-  bool isProductTuple() const { return kind == Kind::ProductTuple; }
+  bool isVector() const { return kind == Kind::Vector; }
+  bool isStruct() const { return kind == Kind::Struct; }
+  bool isTuple() const { return kind == Kind::Tuple; }
 
   Kind getKind() const { return kind; }
-  BuiltinFloatType *getBuiltinRealScalarSpace() const {
-    assert(kind == Kind::BuiltinRealScalar);
+  BuiltinFloatType *getBuiltinFloat() const {
+    assert(kind == Kind::BuiltinFloat);
     return value.builtinFPType;
   }
-  NominalTypeDecl *getRealScalarSpace() const {
-    assert(kind == Kind::RealScalar);
-    return value.realNominalType;
-  }
-  NominalTypeDecl *getRealVectorSpace() const {
-    assert(kind == Kind::RealVector);
-    return value.realNominalType;
-  }
-  NominalTypeDecl *getRealScalarOrVectorSpace() const {
-    assert(kind == Kind::RealScalar || kind == Kind::RealVector);
+  NominalTypeDecl *getVector() const {
+    assert(kind == Kind::Vector);
     return value.realNominalType;
   }
-  StructDecl *getProductStruct() const {
-    assert(kind == Kind::ProductStruct);
+  StructDecl *getStruct() const {
+    assert(kind == Kind::Struct);
     return value.structDecl;
   }
-  TupleType *getProductTuple() const {
-    assert(kind == Kind::ProductTuple);
+  TupleType *getTuple() const {
+    assert(kind == Kind::Tuple);
     return value.tupleType;
   }
-  EnumDecl *getSum() const {
-    assert(kind == Kind::Sum);
+  EnumDecl *getEnum() const {
+    assert(kind == Kind::Enum);
     return value.enumDecl;
   }
 };

lib/AST/ASTContext.cpp

@@ -382,7 +382,7 @@ FOR_KNOWN_FOUNDATION_TYPES(CACHE_FOUNDATION_DECL)
 
   // SWIFT_ENABLE_TENSORFLOW
   /// A cache of tangent spaces per type.
-  llvm::DenseMap<CanType, Optional<TangentSpace>> TangentSpaces;
+  llvm::DenseMap<CanType, Optional<VectorSpace>> VectorSpaces;
 
   /// For uniquifying `AutoDiffParameterIndices` allocations.
   llvm::FoldingSet<AutoDiffParameterIndices> AutoDiffParameterIndicesSet;
@@ -5195,39 +5195,30 @@ LayoutConstraint LayoutConstraint::getLayoutConstraint(LayoutConstraintKind Kind
 }
 
 // SWIFT_ENABLE_TENSORFLOW
-Optional<TangentSpace> ASTContext::getTangentSpace(CanType type,
-                                                   ModuleDecl *module) {
-  auto lookup = getImpl().TangentSpaces.find(type);
-  if (lookup != getImpl().TangentSpaces.end())
+Optional<VectorSpace> ASTContext::getTangentSpace(CanType type,
+                                                  ModuleDecl *module) {
+  auto lookup = getImpl().VectorSpaces.find(type);
+  if (lookup != getImpl().VectorSpaces.end())
     return lookup->getSecond();
   // A helper that is used to cache the computed tangent space for the
   // specified type and retuns the same tangent space.
-  auto cache = [&](Optional<TangentSpace> tangentSpace) {
-    getImpl().TangentSpaces.insert({type, tangentSpace});
-    return tangentSpace;
+  auto cache = [&](Optional<VectorSpace> space) {
+    getImpl().VectorSpaces.insert({type, space});
+    return space;
   };
   // `Builtin.FP<...>` is a builtin real scalar space.
   if (auto *fpType = type->getAs<BuiltinFloatType>())
-    return cache(TangentSpace::getBuiltinRealScalarSpace(fpType));
-  // Look up conformance to `FloatingPoint`.
-  auto *fpProto = getProtocol(KnownProtocolKind::FloatingPoint);
-  if (auto maybeFPConf = module->lookupConformance(type, fpProto)) {
-    auto *typeDecl = type->getAnyNominal();
-    assert(typeDecl);
-    return cache(TangentSpace::getRealScalarSpace(typeDecl));
-  }
+    return cache(VectorSpace::getBuiltinRealScalarSpace(fpType));
   // Look up conformance to `Differentiable`.
   auto *diffableProto = getProtocol(KnownProtocolKind::Differentiable);
-  if (auto maybeDiffableConf = module->lookupConformance(type, diffableProto)) {
-    auto tangentLookup =
-        diffableProto->lookupDirect(getIdentifier("TangentVector"));
-    auto *tangentAssocDecl = cast<AssociatedTypeDecl>(tangentLookup[0]);
-    auto subMap = type->getMemberSubstitutionMap(module, tangentAssocDecl);
-    auto tangent = tangentAssocDecl->getDeclaredInterfaceType().subst(subMap);
-    auto *tangentDecl = tangent->getAnyNominal();
-    assert(tangentDecl &&
-           "Tangent must be a nominal type because it has protocol contraints");
-    return cache(TangentSpace::getRealVectorSpace(tangentDecl));
+  if (module->lookupConformance(type, diffableProto).hasValue()) {
+    auto tangentType = type->getAutoDiffAssociatedType(
+        AutoDiffAssociatedTypeKind::TangentVector,
+        LookUpConformanceInModule(module));
+    assert(tangentType);
+    auto *nomTypeDecl = tangentType->getAnyNominal();
+    assert(nomTypeDecl);
+    return VectorSpace::getRealVectorSpace(nomTypeDecl);
   }
   // Nominal types can be either a struct or an enum.
   if (auto *nominal = type->getAnyNominal()) {
@@ -5243,7 +5234,7 @@ Optional<TangentSpace> ASTContext::getTangentSpace(CanType type,
                                          module);
           });
       if (allMembersHaveTangentSpace)
-        return cache(TangentSpace::getProductStruct(structDecl));
+        return cache(VectorSpace::getStruct(structDecl));
     }
     // Frozen enum types, all of whose payloads have a tangent space, are a
     // sum of the product of payloads in each case.
@@ -5263,15 +5254,15 @@ Optional<TangentSpace> ASTContext::getTangentSpace(CanType type,
           });
         });
       if (allMembersHaveTangentSpace)
-        return cache(TangentSpace::getSum(enumDecl));
+        return cache(VectorSpace::getEnum(enumDecl));
     }
   }
   // Tuple types, each of whose elements has a tangent space, are a product of
   // those tangent space.
   if (TupleType *tupleType = type->getAs<TupleType>())
     if (llvm::all_of(tupleType->getElementTypes(), [&](Type t) {
             return (bool)getTangentSpace(t->getCanonicalType(), module); }))
-      return cache(TangentSpace::getProductTuple(tupleType));
+      return cache(VectorSpace::getTuple(tupleType));
   // Otherwise, the type does not have a tangent space. That is, it does not
   // support differentiation.
   return cache(None);