Merge pull request swiftlang#31173 from dan-zheng/property-wrapper-differentiation

[AutoDiff] Support differentiation of wrapped properties.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -2734,13 +2734,21 @@ 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,
+      (/*propName*/ Identifier, /*propType*/ Type, /*nominalName*/ Identifier,
+       /*nominalCanDeriveAdditiveArithmetic*/ bool))
+WARNING(differentiable_immutable_wrapper_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"
+      "requires all stored properties not marked with `@noDerivative` to be "
+      "mutable; add an explicit '@noDerivative' attribute"
       "%select{|, or conform %1 to 'AdditiveArithmetic'}2",
-      (Identifier, Identifier, bool))
+      (/*wrapperType*/ StringRef, /*nominalName*/ Identifier,
+       /*nominalCanDeriveAdditiveArithmetic*/ bool))
+WARNING(differentiable_let_property_implicit_noderivative_fixit,none,
+      "synthesis of the 'Differentiable.move(along:)' requirement for %0 "
+      "requires all stored properties not marked with `@noDerivative` to be "
+      "mutable; use 'var' instead, or add an explicit '@noDerivative' attribute"
+      "%select{|, or conform %0 to 'AdditiveArithmetic'}1",
+      (/*nominalName*/ Identifier, /*nominalCanDeriveAdditiveArithmetic*/ bool))
 
 NOTE(codable_extraneous_codingkey_case_here,none,
      "CodingKey case %0 does not match any stored properties", (Identifier))

include/swift/SILOptimizer/Differentiation/AdjointValue.h

@@ -165,6 +165,7 @@ class AdjointValue final {
       break;
     }
   }
+  SWIFT_DEBUG_DUMP { print(llvm::dbgs()); };
 };
 
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,

include/swift/SILOptimizer/Differentiation/Common.h

@@ -61,6 +61,23 @@ ApplyInst *getAllocateUninitializedArrayIntrinsicElementAddress(SILValue v);
 /// tuple-typed and such a user exists.
 DestructureTupleInst *getSingleDestructureTupleUser(SILValue value);
 
+/// Returns true if the given original function is a "semantic member accessor".
+///
+/// "Semantic member accessors" are attached to member properties that have a
+/// corresponding tangent stored property in the parent `TangentVector` type.
+/// These accessors have special-case pullback generation based on their
+/// semantic behavior.
+///
+/// "Semantic member accessors" currently include:
+/// - Stored property accessors. These are implicitly generated.
+/// - Property wrapper wrapped value accessors. These are implicitly generated
+///   and internally call `var wrappedValue`.
+bool isSemanticMemberAccessor(SILFunction *original);
+
+/// Returns true if the given apply site has a "semantic member accessor"
+/// callee.
+bool hasSemanticMemberAccessorCallee(ApplySite applySite);
+
 /// Given a full apply site, apply the given callback to each of its
 /// "direct results".
 ///

include/swift/SILOptimizer/Differentiation/PullbackEmitter.h

@@ -100,7 +100,7 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   SILBuilder localAllocBuilder;
 
   /// Stack buffers allocated for storing local adjoint values.
-  SmallVector<SILValue, 64> functionLocalAllocations;
+  SmallVector<AllocStackInst *, 64> functionLocalAllocations;
 
   /// A set used to remember local allocations that were destroyed.
   llvm::SmallDenseSet<SILValue> destroyedLocalAllocations;
@@ -316,6 +316,19 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   /// if any error occurs.
   bool run();
 
+  /// Performs pullback generation on the empty pullback function, given that
+  /// the original function is a "semantic member accessor".
+  ///
+  /// "Semantic member accessors" are attached to member properties that have a
+  /// corresponding tangent stored property in the parent `TangentVector` type.
+  /// These accessors have special-case pullback generation based on their
+  /// semantic behavior.
+  ///
+  /// Returns true if any error occurs.
+  bool runForSemanticMemberAccessor();
+  bool runForSemanticMemberGetter();
+  bool runForSemanticMemberSetter();
+
   /// If original result is non-varied, it will always have a zero derivative.
   /// Skip full pullback generation and simply emit zero derivatives for wrt
   /// parameters.

lib/SILOptimizer/Differentiation/Common.cpp

@@ -71,6 +71,44 @@ DestructureTupleInst *getSingleDestructureTupleUser(SILValue value) {
   return result;
 }
 
+bool isSemanticMemberAccessor(SILFunction *original) {
+  auto *dc = original->getDeclContext();
+  if (!dc)
+    return false;
+  auto *decl = dc->getAsDecl();
+  if (!decl)
+    return false;
+  auto *accessor = dyn_cast<AccessorDecl>(decl);
+  if (!accessor)
+    return false;
+  // Currently, only getters and setters are supported.
+  // TODO(SR-12640): Support `modify` accessors.
+  if (accessor->getAccessorKind() != AccessorKind::Get &&
+      accessor->getAccessorKind() != AccessorKind::Set)
+    return false;
+  // Accessor must come from a `var` declaration.
+  auto *varDecl = dyn_cast<VarDecl>(accessor->getStorage());
+  if (!varDecl)
+    return false;
+  // Return true for stored property accessors.
+  if (varDecl->hasStorage() && varDecl->isInstanceMember())
+    return true;
+  // Return true for properties that have attached property wrappers.
+  if (varDecl->hasAttachedPropertyWrapper())
+    return true;
+  // Otherwise, return false.
+  // User-defined accessors can never be supported because they may use custom
+  // logic that does not semantically perform a member access.
+  return false;
+}
+
+bool hasSemanticMemberAccessorCallee(ApplySite applySite) {
+  if (auto *FRI = dyn_cast<FunctionRefBaseInst>(applySite.getCallee()))
+    if (auto *F = FRI->getReferencedFunctionOrNull())
+      return isSemanticMemberAccessor(F);
+  return false;
+}
+
 void forEachApplyDirectResult(
     FullApplySite applySite,
     llvm::function_ref<void(SILValue)> resultCallback) {