[AutoDiff upstream] Add Differentiable protocol derived conformances. (#30671)

Add `AdditiveArithmetic` derived conformances for structs and classes, gated by
the `-enable-experimental-differentiable-programming` flag.

Structs and classes whose stored properties all conform to `Differentiable` can
derive `Differentiable`:

- `associatedtype TangentVector: Differentiable & AdditiveArithmetic`
  - Member `TangentVector` structs are synthesized whose stored properties are
    all `var` stored properties that conform to `Differentiable` and that are
    not `@noDerivative`.
- `mutating func move(along: TangentVector)`

The `@noDerivative` attribute may be declared on stored properties to opt out of
inclusion in synthesized `TangentVector` structs.

Some stored properties cannot be used in `TangentVector` struct synthesis and
are implicitly marked as `@noDerivative`, with a warning:

- `let` stored properties.
  - These cannot be updated by `mutating func move(along: TangentVector)`.
- Non-`Differentiable`-conforming stored properties.

`@noDerivative` also implies `@_semantics("autodiff.nonvarying")`, which is
relevant for differentiable activity analysis.

Add type-checking and SILGen tests.

Resolves TF-845.

Author: dan-zheng

include/swift/AST/Attr.def

@@ -552,6 +552,11 @@ DECL_ATTR(transpose, Transpose,
   ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove,
   99)
 
+SIMPLE_DECL_ATTR(noDerivative, NoDerivative,
+  OnAbstractFunction | OnVar | OnSubscript |
+  ABIBreakingToAdd | ABIBreakingToRemove | APIBreakingToAdd | APIBreakingToRemove,
+  100)
+
 #undef TYPE_ATTR
 #undef DECL_ATTR_ALIAS
 #undef CONTEXTUAL_DECL_ATTR_ALIAS

include/swift/AST/DiagnosticsSema.def

@@ -2712,6 +2712,19 @@ ERROR(broken_encodable_requirement,none,
       "Encodable protocol is broken: unexpected requirement", ())
 ERROR(broken_decodable_requirement,none,
       "Decodable protocol is broken: unexpected requirement", ())
+ERROR(broken_differentiable_requirement,none,
+      "Differentiable protocol is broken: unexpected requirement", ())
+WARNING(differentiable_nondiff_type_implicit_noderivative_fixit,none,
+      "stored property %0 has no derivative because %1 does not conform to "
+      "'Differentiable'; add an explicit '@noDerivative' attribute"
+      "%select{|, or conform %2 to 'AdditiveArithmetic'}3",
+      (Identifier, Type, Identifier, bool))
+WARNING(differentiable_let_property_implicit_noderivative_fixit,none,
+      "synthesis of the 'Differentiable.move(along:)' requirement for %1 "
+      "requires all stored properties to be mutable; use 'var' instead, or add "
+      "an explicit '@noDerivative' attribute"
+      "%select{|, or conform %1 to 'AdditiveArithmetic'}2",
+      (Identifier, Identifier, bool))
 
 NOTE(codable_extraneous_codingkey_case_here,none,
      "CodingKey case %0 does not match any stored properties", (Identifier))

include/swift/AST/KnownIdentifiers.def

@@ -204,7 +204,10 @@ IDENTIFIER_(nsError)
 IDENTIFIER(OSLogMessage)
 
 // Differentiable programming
+IDENTIFIER(along)
 IDENTIFIER(differential)
+IDENTIFIER(direction)
+IDENTIFIER(move)
 IDENTIFIER(pullback)
 IDENTIFIER(TangentVector)
 IDENTIFIER(zero)

lib/AST/Decl.cpp

@@ -5021,7 +5021,8 @@ ArrayRef<Requirement> ProtocolDecl::getCachedRequirementSignature() const {
 void ProtocolDecl::computeKnownProtocolKind() const {
   auto module = getModuleContext();
   if (module != module->getASTContext().getStdlibModule() &&
-      !module->getName().is("Foundation")) {
+      !module->getName().is("Foundation") &&
+      !module->getName().is("_Differentiation")) {
     const_cast<ProtocolDecl *>(this)->Bits.ProtocolDecl.KnownProtocol = 1;
     return;
   }

lib/SIL/SILFunctionBuilder.cpp

@@ -66,6 +66,25 @@ void SILFunctionBuilder::addFunctionAttributes(
   if (Attrs.hasAttribute<SILGenNameAttr>() || Attrs.hasAttribute<CDeclAttr>())
     F->setHasCReferences(true);
 
+  // Validate `@differentiable` attributes by calling `getParameterIndices`.
+  // This is important for:
+  // - Skipping invalid `@differentiable` attributes in non-primary files.
+  // - Preventing duplicate SIL differentiability witness creation for
+  //   `@differentiable` attributes on `AbstractStorageDecl` declarations.
+  //   Such `@differentiable` attributes are deleted and recreated on the getter
+  //   `AccessorDecl` of the `AbstractStorageDecl`.
+  for (auto *A : Attrs.getAttributes<DifferentiableAttr>())
+    (void)A->getParameterIndices();
+
+  // Propagate `@noDerivative` as `[_semantics "autodiff.nonvarying"]`.
+  //
+  // `@noDerivative` implies non-varying semantics for differentiable activity
+  // analysis. SIL values produced from references to `@noDerivative`
+  // declarations will not be marked as varying; these values do not need a
+  // derivative.
+  if (Attrs.hasAttribute<NoDerivativeAttr>())
+    F->addSemanticsAttr("autodiff.nonvarying");
+
   // Propagate @_dynamicReplacement(for:).
   if (constant.isNull())
     return;

lib/Sema/CMakeLists.txt

@@ -19,6 +19,7 @@ add_swift_host_library(swiftSema STATIC
   DerivedConformanceCaseIterable.cpp
   DerivedConformanceCodable.cpp
   DerivedConformanceCodingKey.cpp
+  DerivedConformanceDifferentiable.cpp
   DerivedConformanceEquatableHashable.cpp
   DerivedConformanceComparable.cpp
   DerivedConformanceError.cpp

lib/Sema/CodeSynthesis.cpp

@@ -1369,3 +1369,21 @@ bool swift::hasLetStoredPropertyWithInitialValue(NominalTypeDecl *nominal) {
     return v->isLet() && v->hasInitialValue();
   });
 }
+
+void swift::addFixedLayoutAttr(NominalTypeDecl *nominal) {
+  auto &C = nominal->getASTContext();
+  // If nominal already has `@_fixed_layout`, return.
+  if (nominal->getAttrs().hasAttribute<FixedLayoutAttr>())
+    return;
+  auto access = nominal->getEffectiveAccess();
+  // If nominal does not have at least internal access, return.
+  if (access < AccessLevel::Internal)
+    return;
+  // If nominal is internal, it should have the `@usableFromInline` attribute.
+  if (access == AccessLevel::Internal &&
+      !nominal->getAttrs().hasAttribute<UsableFromInlineAttr>()) {
+    nominal->getAttrs().add(new (C) UsableFromInlineAttr(/*Implicit*/ true));
+  }
+  // Add `@_fixed_layout` to the nominal.
+  nominal->getAttrs().add(new (C) FixedLayoutAttr(/*Implicit*/ true));
+}

lib/Sema/CodeSynthesis.h

@@ -67,6 +67,9 @@ ValueDecl *getProtocolRequirement(ProtocolDecl *protocol, Identifier name);
 // with an initial value.
 bool hasLetStoredPropertyWithInitialValue(NominalTypeDecl *nominal);
 
+/// Add `@_fixed_layout` attribute to the nominal type, if possible.
+void addFixedLayoutAttr(NominalTypeDecl *nominal);
+
 } // end namespace swift
 
 #endif