[AutoDiff] Finish wrapped property differentiation support.

Add special-case VJP generation support for "semantic member accessors".
This is necessary to avoid activity analysis related diagnostics and simplifies
generated code.

Fix "wrapped property mutability" check in `Differentiable` derived conformnances.
This resolves SR-12642.

Add e2e test using real world property wrappers (`@Lazy` and `@Clamping`).

Author: dan-zheng

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

@@ -324,11 +324,6 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
   /// 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`.
-  ///
   /// Returns true if any error occurs.
   bool runForSemanticMemberAccessor();
   bool runForSemanticMemberGetter();

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) {

lib/SILOptimizer/Differentiation/PullbackEmitter.cpp

@@ -517,40 +517,6 @@ void PullbackEmitter::addToAdjointBuffer(SILBasicBlock *origBB,
 // Member accessor pullback generation
 //--------------------------------------------------------------------------//
 
-/// Returns true if the given original function is a "semantic member accessor".
-static 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())
-    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.
-  // TODO(SR-12636): Support `@differentiable(useInTangentVector)` computed
-  // properties.
-  return false;
-}
-
 bool PullbackEmitter::runForSemanticMemberAccessor() {
   auto &original = getOriginal();
   auto *accessor = cast<AccessorDecl>(original.getDeclContext()->getAsDecl());
@@ -593,7 +559,8 @@ bool PullbackEmitter::runForSemanticMemberGetter() {
   auto origResult = origFormalResults[getIndices().source];
 
   // TODO(TF-970): Emit diagnostic when `TangentVector` is not a struct.
-  auto tangentVectorSILTy = pullback.getConventions().getSingleSILResultType();
+  auto tangentVectorSILTy = pullback.getConventions().getSingleSILResultType(
+      TypeExpansionContext::minimal());
   auto tangentVectorTy = tangentVectorSILTy.getASTType();
   auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
 
@@ -806,15 +773,6 @@ bool PullbackEmitter::run() {
         return;
       visited.insert(v);
       auto type = v->getType();
-      // Diagnose active enum values. Differentiation of enum values requires
-      // special adjoint value handling and is not yet supported. Diagnose
-      // only the first active enum value to prevent too many diagnostics.
-      if (!diagnosedActiveEnumValue && type.getEnumOrBoundGenericEnum()) {
-        getContext().emitNondifferentiabilityError(
-            v, getInvoker(), diag::autodiff_enums_unsupported);
-        errorOccurred = true;
-        diagnosedActiveEnumValue = true;
-      }
       // Diagnose active values whose value category is incompatible with their
       // tangent types's value category.
       //
@@ -842,6 +800,19 @@ bool PullbackEmitter::run() {
           }
         }
       }
+      // Do not emit remaining activity-related diagnostics for semantic member
+      // accessors, which have special-case pullback generation.
+      if (isSemanticMemberAccessor(&original))
+        return;
+      // Diagnose active enum values. Differentiation of enum values requires
+      // special adjoint value handling and is not yet supported. Diagnose
+      // only the first active enum value to prevent too many diagnostics.
+      if (!diagnosedActiveEnumValue && type.getEnumOrBoundGenericEnum()) {
+        getContext().emitNondifferentiabilityError(
+            v, getInvoker(), diag::autodiff_enums_unsupported);
+        errorOccurred = true;
+        diagnosedActiveEnumValue = true;
+      }
       // Skip address projections.
       // Address projections do not need their own adjoint buffers; they
       // become projections into their adjoint base buffer.

lib/SILOptimizer/Differentiation/VJPEmitter.cpp

@@ -511,14 +511,35 @@ void VJPEmitter::visitSwitchEnumAddrInst(SwitchEnumAddrInst *seai) {
 }
 
 void VJPEmitter::visitApplyInst(ApplyInst *ai) {
-  // If the function should not be differentiated or its the array literal
-  // initialization intrinsic, just do standard cloning.
-  if (!pullbackInfo.shouldDifferentiateApplySite(ai) ||
-      isArrayLiteralIntrinsic(ai)) {
+  // If callee should not be differentiated, do standard cloning.
+  if (!pullbackInfo.shouldDifferentiateApplySite(ai)) {
     LLVM_DEBUG(getADDebugStream() << "No active results:\n" << *ai << '\n');
     TypeSubstCloner::visitApplyInst(ai);
     return;
   }
+  // If callee is the array literal initialization intrinsic, do standard
+  // cloning. Array literal differentiation is handled separately.
+  if (isArrayLiteralIntrinsic(ai)) {
+    LLVM_DEBUG(getADDebugStream() << "Cloning array literal intrinsic `apply`\n"
+                                  << *ai << '\n');
+    TypeSubstCloner::visitApplyInst(ai);
+    return;
+  }
+  // If the original function is a semantic member accessor, do standard
+  // cloning. Semantic member accessors have special pullback generation logic,
+  // so all `apply` instructions can be directly cloned to the VJP.
+  if (isSemanticMemberAccessor(original)) {
+    LLVM_DEBUG(getADDebugStream()
+               << "Cloning `apply` in semantic member accessor:\n"
+               << *ai << '\n');
+    TypeSubstCloner::visitApplyInst(ai);
+    return;
+  }
+
+  auto loc = ai->getLoc();
+  auto &builder = getBuilder();
+  auto origCallee = getOpValue(ai->getCallee());
+  auto originalFnTy = origCallee->getType().castTo<SILFunctionType>();
 
   LLVM_DEBUG(getADDebugStream() << "VJP-transforming:\n" << *ai << '\n');
 
@@ -558,31 +579,27 @@ void VJPEmitter::visitApplyInst(ApplyInst *ai) {
                        activeParamIndices));
 
   // Emit the VJP.
-  auto loc = ai->getLoc();
-  auto &builder = getBuilder();
-  auto original = getOpValue(ai->getCallee());
   SILValue vjpValue;
   // If functionSource is a `@differentiable` function, just extract it.
-  auto originalFnTy = original->getType().castTo<SILFunctionType>();
   if (originalFnTy->isDifferentiable()) {
     auto paramIndices = originalFnTy->getDifferentiabilityParameterIndices();
     for (auto i : indices.parameters->getIndices()) {
       if (!paramIndices->contains(i)) {
         context.emitNondifferentiabilityError(
-            original, invoker,
+            origCallee, invoker,
             diag::autodiff_function_noderivative_parameter_not_differentiable);
         errorOccurred = true;
         return;
       }
     }
-    auto origFnType = original->getType().castTo<SILFunctionType>();
+    auto origFnType = origCallee->getType().castTo<SILFunctionType>();
     auto origFnUnsubstType = origFnType->getUnsubstitutedType(getModule());
     if (origFnType != origFnUnsubstType) {
-      original = builder.createConvertFunction(
-          loc, original, SILType::getPrimitiveObjectType(origFnUnsubstType),
+      origCallee = builder.createConvertFunction(
+          loc, origCallee, SILType::getPrimitiveObjectType(origFnUnsubstType),
           /*withoutActuallyEscaping*/ false);
     }
-    auto borrowedDiffFunc = builder.emitBeginBorrowOperation(loc, original);
+    auto borrowedDiffFunc = builder.emitBeginBorrowOperation(loc, origCallee);
     vjpValue = builder.createDifferentiableFunctionExtract(
         loc, NormalDifferentiableFunctionTypeComponent::VJP, borrowedDiffFunc);
     vjpValue = builder.emitCopyValueOperation(loc, vjpValue);
@@ -624,7 +641,7 @@ void VJPEmitter::visitApplyInst(ApplyInst *ai) {
         }
         if (!remappedResultType.isDifferentiable(getModule())) {
           context.emitNondifferentiabilityError(
-              original, invoker, diag::autodiff_nondifferentiable_result);
+              origCallee, invoker, diag::autodiff_nondifferentiable_result);
           errorOccurred = true;
           return true;
         }
@@ -645,7 +662,7 @@ void VJPEmitter::visitApplyInst(ApplyInst *ai) {
     /*
     DifferentiationInvoker indirect(ai, attr);
     auto insertion =
-        context.getInvokers().try_emplace({this->original, attr}, indirect);
+        context.getInvokers().try_emplace({original, attr}, indirect);
     auto &invoker = insertion.first->getSecond();
     invoker = indirect;
     */
@@ -656,21 +673,21 @@ void VJPEmitter::visitApplyInst(ApplyInst *ai) {
     // function operand is specialized with a remapped version of same
     // substitution map using an argument-less `partial_apply`.
     if (ai->getSubstitutionMap().empty()) {
-      original = builder.emitCopyValueOperation(loc, original);
+      origCallee = builder.emitCopyValueOperation(loc, origCallee);
     } else {
       auto substMap = getOpSubstitutionMap(ai->getSubstitutionMap());
       auto vjpPartialApply = getBuilder().createPartialApply(
-          ai->getLoc(), original, substMap, {},
+          ai->getLoc(), origCallee, substMap, {},
           ParameterConvention::Direct_Guaranteed);
-      original = vjpPartialApply;
-      originalFnTy = original->getType().castTo<SILFunctionType>();
+      origCallee = vjpPartialApply;
+      originalFnTy = origCallee->getType().castTo<SILFunctionType>();
       // Diagnose if new original function type is non-differentiable.
       if (diagnoseNondifferentiableOriginalFunctionType(originalFnTy))
         return;
     }
 
     auto *diffFuncInst = context.createDifferentiableFunction(
-        getBuilder(), loc, indices.parameters, original);
+        getBuilder(), loc, indices.parameters, origCallee);
 
     // Record the `differentiable_function` instruction.
     context.addDifferentiableFunctionInstToWorklist(diffFuncInst);

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -42,11 +42,11 @@ getStoredPropertiesForDifferentiation(NominalTypeDecl *nominal, DeclContext *DC,
   for (auto *vd : nominal->getStoredProperties()) {
     // Peer through property wrappers: use original wrapped properties instead.
     if (auto *originalProperty = vd->getOriginalWrappedProperty()) {
-      // Skip wrapped properties that do not define a setter. They define a setter
-      // only if property wrappers define a setter for `var wrappedValue`.
-      // `mutating func move(along:)` cannot be synthesized to update these
-      // properties.
-      if (!originalProperty->getAccessor(AccessorKind::Set))
+      // Skip immutable wrapped properties. `mutating func move(along:)` cannot
+      // be synthesized to update these properties.
+      auto mutability = originalProperty->getPropertyWrapperMutability();
+      assert(mutability.hasValue() && "Expected wrapped property mutability");
+      if (mutability->Setter != PropertyWrapperMutability::Value::Mutating)
         continue;
       // Use the original wrapped property.
       vd = originalProperty;
@@ -510,9 +510,9 @@ static void checkAndDiagnoseImplicitNoDerivative(ASTContext &Context,
       // to update these properties.
       auto *wrapperDecl =
           vd->getInterfaceType()->getNominalOrBoundGenericNominal();
-      auto *wrappedValueDecl =
-          wrapperDecl->getPropertyWrapperTypeInfo().valueVar;
-      if (!wrappedValueDecl->getAccessor(AccessorKind::Set)) {
+      auto mutability = originalProperty->getPropertyWrapperMutability();
+      assert(mutability.hasValue() && "Expected wrapped property mutability");
+      if (mutability->Setter != PropertyWrapperMutability::Value::Mutating) {
         auto loc =
             originalProperty->getAttributeInsertionLoc(/*forModifier*/ false);
         Context.Diags

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -177,14 +177,30 @@ let _: @differentiable (Float) -> TF_691<Float> = { x in id(TF_691(x)) }
 
 @propertyWrapper
 struct Wrapper<Value> {
-  var wrappedValue: Value
+  private var value: Value
+  var wrappedValue: Value {
+    get { value }
+    set { value = newValue }
+  }
   var projectedValue: Self { self }
+
+  init(wrappedValue: Value) {
+    self.value = wrappedValue
+  }
 }
 
 @propertyWrapper
 struct DifferentiableWrapper<Value> {
-  var wrappedValue: Value
+  private var value: Value
+  var wrappedValue: Value {
+    get { value }
+    set { value = newValue }
+  }
   var projectedValue: Self { self }
+
+  init(wrappedValue: Value) {
+    self.value = wrappedValue
+  }
 }
 extension DifferentiableWrapper: Differentiable where Value: Differentiable {}
 // Note: property wrapped value differentiation works even if wrapper types do

test/AutoDiff/compiler_crashers/sr12642-differentiable-derivation-redeclared-property.swift

@@ -0,0 +1,25 @@
+// RUN: %target-swift-frontend -typecheck -verify %s
+// REQUIRES: asserts
+
+// SR-12642: Crash regarding `Differentiable` derived conformances and
+// redeclared properties. This crash surfaced only briefly during the
+// implementation of wrapped property differentiation (SR-12637).
+
+import _Differentiation
+
+@propertyWrapper
+struct Wrapper<Value> {
+  var wrappedValue: Value
+}
+
+struct Generic<T> {}
+extension Generic: Differentiable where T: Differentiable {}
+
+struct WrappedProperties: Differentiable {
+  // expected-note @+2 {{'int' previously declared here}}
+  // expected-warning @+1 {{stored property 'int' has no derivative because 'Generic<Int>' does not conform to 'Differentiable'; add an explicit '@noDerivative' attribute}}
+  @Wrapper var int: Generic<Int>
+
+  // expected-error @+1 {{invalid redeclaration of 'int'}}
+  @Wrapper var int: Generic<Int>
+}