Derive ElementaryFunctions conformances for structs. (#25500)

`ElementaryFunctions` derived conformances enable elementary math
functions to work with product spaces formed from
`ElementaryFunctions`-conforming types.

Conform `Differentiable` synthesized associated types to
`ElementaryFunctions` if possible. Similar to existing logic for `VectorProtocol`.

Enables efficient, elegant mathematical optimizers.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -2532,6 +2532,8 @@ ERROR(parameterized_invalid_parameters_struct,none,
       "invalid", (Type))
 ERROR(broken_additive_arithmetic_requirement,none,
       "AdditiveArithmetic protocol is broken: unexpected requirement", ())
+ERROR(broken_elementary_functions_requirement,none,
+      "ElementaryFunctions protocol is broken: unexpected requirement", ())
 ERROR(broken_vector_protocol_requirement,none,
       "VectorProtocol protocol is broken: unexpected requirement", ())
 ERROR(broken_differentiable_requirement,none,

include/swift/AST/KnownProtocols.def

@@ -78,6 +78,7 @@ PROTOCOL(Encodable)
 PROTOCOL(Decodable)
 // SWIFT_ENABLE_TENSORFLOW
 PROTOCOL(AdditiveArithmetic)
+PROTOCOL(ElementaryFunctions)
 PROTOCOL(KeyPathIterable)
 PROTOCOL(TensorArrayProtocol)
 PROTOCOL(TensorGroup)

lib/IRGen/GenMeta.cpp

@@ -4206,6 +4206,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::StringInterpolationProtocol:
   // SWIFT_ENABLE_TENSORFLOW
   case KnownProtocolKind::AdditiveArithmetic:
+  case KnownProtocolKind::ElementaryFunctions:
   case KnownProtocolKind::KeyPathIterable:
   case KnownProtocolKind::TensorArrayProtocol:
   case KnownProtocolKind::TensorGroup:

lib/Sema/CMakeLists.txt

@@ -28,6 +28,7 @@ add_swift_host_library(swiftSema STATIC
   DerivedConformanceError.cpp
   # SWIFT_ENABLE_TENSORFLOW
   DerivedConformanceAdditiveArithmetic.cpp
+  DerivedConformanceElementaryFunctions.cpp
   DerivedConformanceVectorProtocol.cpp
   DerivedConformanceDifferentiable.cpp
   DerivedConformanceKeyPathIterable.cpp

lib/Sema/DerivedConformanceAdditiveArithmetic.cpp

@@ -183,7 +183,7 @@ static void deriveBodyMathOperator(AbstractFunctionDecl *funcDecl,
     memberOpExprs.push_back(createMemberOpExpr(member));
     memberNames.push_back(member->getName());
   }
-  // Call memberwise initialier with member operator call expressions.
+  // Call memberwise initializer with member operator call expressions.
   auto *callExpr =
       CallExpr::createImplicit(C, initExpr, memberOpExprs, memberNames);
   ASTNode returnStmt = new (C) ReturnStmt(SourceLoc(), callExpr, true);

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -568,12 +568,14 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
   auto diffableType = TypeLoc::withoutLoc(diffableProto->getDeclaredType());
   auto *addArithProto = C.getProtocol(KnownProtocolKind::AdditiveArithmetic);
   auto addArithType = TypeLoc::withoutLoc(addArithProto->getDeclaredType());
+  auto *mathProto = C.getProtocol(KnownProtocolKind::ElementaryFunctions);
+  auto mathType = TypeLoc::withoutLoc(mathProto->getDeclaredType());
   auto *vectorProto = C.getProtocol(KnownProtocolKind::VectorProtocol);
   auto vectorType = TypeLoc::withoutLoc(vectorProto->getDeclaredType());
   auto *kpIterableProto = C.getProtocol(KnownProtocolKind::KeyPathIterable);
   auto kpIterableType = TypeLoc::withoutLoc(kpIterableProto->getDeclaredType());
 
-  SmallVector<TypeLoc, 3> inherited {diffableType};
+  SmallVector<TypeLoc, 3> inherited{diffableType};
 
   // Cache original members and their associated types for later use.
   SmallVector<VarDecl *, 8> diffProperties;
@@ -589,9 +591,16 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
   bool canDeriveAdditiveArithmetic =
       llvm::all_of(diffProperties, [&](VarDecl *vd) {
         return TC.conformsToProtocol(getAssociatedType(vd, parentDC, id),
-                                     addArithProto, parentDC,
-                                     None);
-        });
+                                     addArithProto, parentDC, None);
+      });
+
+  // Associated struct can derive `ElementaryFunctions` if the associated types
+  // of all stored properties conform to `ElementaryFunctions`.
+  bool canDeriveElementaryFunctions =
+      llvm::all_of(diffProperties, [&](VarDecl *vd) {
+        return TC.conformsToProtocol(getAssociatedType(vd, parentDC, id),
+                                     mathProto, parentDC, None);
+      });
 
   // Associated struct can derive `VectorProtocol` if the associated types of
   // all members conform to `VectorProtocol` and share the same
@@ -625,6 +634,10 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
                               None))
       inherited.push_back(kpIterableType);
   }
+  // If all members conform to `ElementaryFunctions`, make the associated struct
+  // conform to `ElementaryFunctions`.
+  if (canDeriveElementaryFunctions)
+    inherited.push_back(mathType);
   // If all members also conform to `VectorProtocol` with the same `Scalar`
   // type, make the associated struct conform to `VectorProtocol` instead of
   // just `AdditiveArithmetic`.