[AutoDiff] Support property wrapper differentiation.

Support differentiation of property wrapper wrapped value getters and setters.

Create new pullback generation code path for "semantic member accessors".
"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. Currently, only getters and setters are supported.

This special-case pullback generation is currently used for stored property
accessors and property wrapper wrapped value accessors. In the future, it can
also be used to support `@differentiable(useInTangentVector)` computed
properties: SR-12636.

User-defined accesors cannot use this code path because they may use custom
logic that does not semantically perform a member access.

Resolves SR-12639.

Author: dan-zheng

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/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,24 @@ 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.
+  ///
+  /// "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();
+  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/PullbackEmitter.cpp

@@ -22,6 +22,8 @@
 #include "swift/SILOptimizer/Differentiation/Thunk.h"
 #include "swift/SILOptimizer/Differentiation/VJPEmitter.h"
 
+#include "swift/AST/Expr.h"
+#include "swift/AST/PropertyWrappers.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SILOptimizer/PassManager/PrettyStackTrace.h"
@@ -511,6 +513,242 @@ void PullbackEmitter::addToAdjointBuffer(SILBasicBlock *origBB,
   accumulateIndirect(adjointBuffer, rhsBufferAccess, loc);
 }
 
+//--------------------------------------------------------------------------//
+// 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());
+  switch (accessor->getAccessorKind()) {
+  case AccessorKind::Get:
+    return runForSemanticMemberGetter();
+  case AccessorKind::Set:
+    return runForSemanticMemberSetter();
+  // TODO(SR-12640): Support `modify` accessors.
+  default:
+    llvm_unreachable("Unsupported accessor kind; inconsistent with "
+                     "`isSemanticMemberAccessor`?");
+  }
+}
+
+bool PullbackEmitter::runForSemanticMemberGetter() {
+  auto &original = getOriginal();
+  auto &pullback = getPullback();
+  auto pbLoc = getPullback().getLocation();
+
+  auto *accessor = cast<AccessorDecl>(original.getDeclContext()->getAsDecl());
+  assert(accessor->getAccessorKind() == AccessorKind::Get);
+
+  auto origExitIt = original.findReturnBB();
+  assert(origExitIt != original.end() &&
+         "Functions without returns must have been diagnosed");
+  auto *origExit = &*origExitIt;
+  builder.setInsertionPoint(pullback.getEntryBlock());
+
+  // Get getter argument and result values.
+  //   Getter type: $(Self) -> Result
+  // Pullback type: $(Result', PB_Struct) -> Self'
+  assert(original.getLoweredFunctionType()->getNumParameters() == 1);
+  assert(pullback.getLoweredFunctionType()->getNumParameters() == 2);
+  assert(pullback.getLoweredFunctionType()->getNumResults() == 1);
+  SILValue origSelf = original.getArgumentsWithoutIndirectResults().front();
+
+  SmallVector<SILValue, 8> origFormalResults;
+  collectAllFormalResultsInTypeOrder(original, origFormalResults);
+  auto origResult = origFormalResults[getIndices().source];
+
+  // TODO(TF-970): Emit diagnostic when `TangentVector` is not a struct.
+  auto tangentVectorSILTy = pullback.getConventions().getSingleSILResultType();
+  auto tangentVectorTy = tangentVectorSILTy.getASTType();
+  auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
+
+  // Look up the corresponding field in the tangent space.
+  VarDecl *origField = cast<VarDecl>(accessor->getStorage());
+  VarDecl *tanField = nullptr;
+  auto tanFieldLookup = tangentVectorDecl->lookupDirect(origField->getName());
+  if (tanFieldLookup.empty()) {
+    getContext().emitNondifferentiabilityError(
+        pbLoc.getSourceLoc(), getInvoker(),
+        diag::autodiff_stored_property_no_corresponding_tangent,
+        origSelf->getType().getASTType().getString(), origField->getNameStr());
+    errorOccurred = true;
+    return true;
+  }
+  tanField = cast<VarDecl>(tanFieldLookup.front());
+
+  // Switch based on the base tangent struct's value category.
+  // TODO(TF-1255): Simplify using unified adjoint value data structure.
+  switch (tangentVectorSILTy.getCategory()) {
+  case SILValueCategory::Object: {
+    auto adjResult = getAdjointValue(origExit, origResult);
+    switch (adjResult.getKind()) {
+    case AdjointValueKind::Zero:
+      addAdjointValue(origExit, origSelf,
+                      makeZeroAdjointValue(tangentVectorSILTy), pbLoc);
+      break;
+    case AdjointValueKind::Concrete:
+    case AdjointValueKind::Aggregate: {
+      SmallVector<AdjointValue, 8> eltVals;
+      for (auto *field : tangentVectorDecl->getStoredProperties()) {
+        if (field == tanField) {
+          eltVals.push_back(adjResult);
+        } else {
+          auto substMap = tangentVectorTy->getMemberSubstitutionMap(
+              field->getModuleContext(), field);
+          auto fieldTy = field->getType().subst(substMap);
+          auto fieldSILTy = getTypeLowering(fieldTy).getLoweredType();
+          assert(fieldSILTy.isObject());
+          eltVals.push_back(makeZeroAdjointValue(fieldSILTy));
+        }
+      }
+      addAdjointValue(origExit, origSelf,
+                      makeAggregateAdjointValue(tangentVectorSILTy, eltVals),
+                      pbLoc);
+    }
+    }
+    break;
+  }
+  case SILValueCategory::Address: {
+    assert(pullback.getIndirectResults().size() == 1);
+    auto pbIndRes = pullback.getIndirectResults().front();
+    auto *adjSelf = createFunctionLocalAllocation(
+        pbIndRes->getType().getObjectType(), pbLoc);
+    setAdjointBuffer(origExit, origSelf, adjSelf);
+    for (auto *field : tangentVectorDecl->getStoredProperties()) {
+      auto *adjSelfElt = builder.createStructElementAddr(pbLoc, adjSelf, field);
+      if (field == tanField) {
+        // Switch based on the property's value category.
+        // TODO(TF-1255): Simplify using unified adjoint value data structure.
+        switch (origResult->getType().getCategory()) {
+        case SILValueCategory::Object: {
+          auto adjResult = getAdjointValue(origExit, origResult);
+          auto adjResultValue = materializeAdjointDirect(adjResult, pbLoc);
+          builder.emitStoreValueOperation(pbLoc, adjResultValue, adjSelfElt,
+                                          StoreOwnershipQualifier::Init);
+          break;
+        }
+        case SILValueCategory::Address: {
+          auto adjResult = getAdjointBuffer(origExit, origResult);
+          builder.createCopyAddr(pbLoc, adjResult, adjSelfElt, IsTake,
+                                 IsInitialization);
+          destroyedLocalAllocations.insert(adjResult);
+          break;
+        }
+        }
+      } else {
+        auto fieldType = pullback.mapTypeIntoContext(field->getInterfaceType())
+                             ->getCanonicalType();
+        emitZeroIndirect(fieldType, adjSelfElt, pbLoc);
+      }
+    }
+    break;
+  }
+  }
+  return false;
+}
+
+bool PullbackEmitter::runForSemanticMemberSetter() {
+  auto &original = getOriginal();
+  auto &pullback = getPullback();
+  auto pbLoc = getPullback().getLocation();
+
+  auto *accessor = cast<AccessorDecl>(original.getDeclContext()->getAsDecl());
+  assert(accessor->getAccessorKind() == AccessorKind::Set);
+
+  auto origEntry = original.getEntryBlock();
+  builder.setInsertionPoint(pullback.getEntryBlock());
+
+  // Get setter argument values.
+  //              Setter type: $(inout Self, Argument) -> ()
+  // Pullback type (wrt self): $(inout Self', PB_Struct) -> ()
+  // Pullback type (wrt both): $(inout Self', PB_Struct) -> Argument'
+  assert(original.getLoweredFunctionType()->getNumParameters() == 2);
+  assert(pullback.getLoweredFunctionType()->getNumParameters() == 2);
+  assert(pullback.getLoweredFunctionType()->getNumResults() == 0 ||
+         pullback.getLoweredFunctionType()->getNumResults() == 1);
+
+  SILValue origArg = original.getArgumentsWithoutIndirectResults()[0];
+  SILValue origSelf = original.getArgumentsWithoutIndirectResults()[1];
+
+  // TODO(TF-970): Emit diagnostic when `TangentVector` is not a struct.
+  auto tangentVectorSILTy = pullback.getLoweredFunctionType()
+                                ->getParameters()[0]
+                                .getSILStorageInterfaceType();
+  assert(tangentVectorSILTy.getCategory() == SILValueCategory::Address);
+  auto tangentVectorTy = tangentVectorSILTy.getASTType();
+  auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
+
+  // Look up the corresponding field in the tangent space.
+  VarDecl *origField = cast<VarDecl>(accessor->getStorage());
+  VarDecl *tanField = nullptr;
+  auto tanFieldLookup = tangentVectorDecl->lookupDirect(origField->getName());
+  if (tanFieldLookup.empty()) {
+    getContext().emitNondifferentiabilityError(
+        pbLoc.getSourceLoc(), getInvoker(),
+        diag::autodiff_stored_property_no_corresponding_tangent,
+        origSelf->getType().getASTType().getString(), origField->getNameStr());
+    errorOccurred = true;
+    return true;
+  }
+  tanField = cast<VarDecl>(tanFieldLookup.front());
+
+  auto adjSelf = getAdjointBuffer(origEntry, origSelf);
+  auto *adjSelfElt = builder.createStructElementAddr(pbLoc, adjSelf, tanField);
+  // Switch based on the property's value category.
+  // TODO(TF-1255): Simplify using unified adjoint value data structure.
+  switch (origArg->getType().getCategory()) {
+  case SILValueCategory::Object: {
+    auto adjArg = builder.emitLoadValueOperation(pbLoc, adjSelfElt,
+                                                 LoadOwnershipQualifier::Take);
+    setAdjointValue(origEntry, origArg, makeConcreteAdjointValue(adjArg));
+    break;
+  }
+  case SILValueCategory::Address: {
+    addToAdjointBuffer(origEntry, origArg, adjSelfElt, pbLoc);
+    builder.emitDestroyOperation(pbLoc, adjSelfElt);
+    break;
+  }
+  }
+  emitZeroIndirect(adjSelfElt->getType().getASTType(), adjSelfElt, pbLoc);
+
+  return false;
+}
+
 //--------------------------------------------------------------------------//
 // Entry point
 //--------------------------------------------------------------------------//
@@ -756,12 +994,21 @@ bool PullbackEmitter::run() {
                << " as the adjoint of original result " << origResult);
   }
 
+  // If the original function is an accessor with special-case pullback
+  // generation logic, do special-case generation.
+  if (isSemanticMemberAccessor(&original)) {
+    if (runForSemanticMemberAccessor())
+      return true;
+  }
+  // Otherwise, perform standard pullback generation.
   // Visit original blocks blocks in post-order and perform differentiation
   // in corresponding pullback blocks. If errors occurred, back out.
-  for (auto *bb : postOrderPostDomOrder) {
-    visitSILBasicBlock(bb);
-    if (errorOccurred)
-      return true;
+  else {
+    for (auto *bb : postOrderPostDomOrder) {
+      visitSILBasicBlock(bb);
+      if (errorOccurred)
+        return true;
+    }
   }
 
   // Prepare and emit a `return` in the pullback exit block.
@@ -1343,8 +1590,7 @@ void PullbackEmitter::visitStructInst(StructInst *si) {
     auto adjStruct = materializeAdjointDirect(std::move(av), loc);
     // Find the struct `TangentVector` type.
     auto structTy = remapType(si->getType()).getASTType();
-    auto tangentVectorTy =
-        getTangentSpace(structTy)->getType()->getCanonicalType();
+    auto tangentVectorTy = getTangentSpace(structTy)->getCanonicalType();
     assert(!getTypeLowering(tangentVectorTy).isAddressOnly());
     auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
     assert(tangentVectorDecl);
@@ -1406,8 +1652,7 @@ void PullbackEmitter::visitStructExtractInst(StructExtractInst *sei) {
          "differentiated; activity analysis should not marked as varied");
   auto *bb = sei->getParent();
   auto structTy = remapType(sei->getOperand()->getType()).getASTType();
-  auto tangentVectorTy =
-      getTangentSpace(structTy)->getType()->getCanonicalType();
+  auto tangentVectorTy = getTangentSpace(structTy)->getCanonicalType();
   assert(!getTypeLowering(tangentVectorTy).isAddressOnly());
   auto tangentVectorSILTy = SILType::getPrimitiveObjectType(tangentVectorTy);
   auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
@@ -1449,8 +1694,7 @@ void PullbackEmitter::visitStructExtractInst(StructExtractInst *sei) {
         auto substMap = tangentVectorTy->getMemberSubstitutionMap(
             field->getModuleContext(), field);
         auto fieldTy = field->getType().subst(substMap);
-        auto fieldSILTy = getContext().getTypeConverter().getLoweredType(
-            fieldTy, TypeExpansionContext::minimal());
+        auto fieldSILTy = getTypeLowering(fieldTy).getLoweredType();
         assert(fieldSILTy.isObject());
         eltVals.push_back(makeZeroAdjointValue(fieldSILTy));
       }
@@ -1466,8 +1710,7 @@ void PullbackEmitter::visitRefElementAddrInst(RefElementAddrInst *reai) {
   auto *bb = reai->getParent();
   auto adjBuf = getAdjointBuffer(bb, reai);
   auto classTy = remapType(reai->getOperand()->getType()).getASTType();
-  auto tangentVectorTy =
-      getTangentSpace(classTy)->getType()->getCanonicalType();
+  auto tangentVectorTy = getTangentSpace(classTy)->getCanonicalType();
   assert(!getTypeLowering(tangentVectorTy).isAddressOnly());
   auto tangentVectorSILTy = SILType::getPrimitiveObjectType(tangentVectorTy);
   auto *tangentVectorDecl = tangentVectorTy->getStructOrBoundGenericStruct();
@@ -1498,8 +1741,7 @@ void PullbackEmitter::visitRefElementAddrInst(RefElementAddrInst *reai) {
       auto substMap = tangentVectorTy->getMemberSubstitutionMap(
           field->getModuleContext(), field);
       auto fieldTy = field->getType().subst(substMap);
-      auto fieldSILTy = getContext().getTypeConverter().getLoweredType(
-          fieldTy, TypeExpansionContext::minimal());
+      auto fieldSILTy = getTypeLowering(fieldTy).getLoweredType();
       assert(fieldSILTy.isObject());
       eltVals.push_back(makeZeroAdjointValue(fieldSILTy));
     }

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -42,14 +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 property wrappers that do not define `wrappedValue.set`.
+      // 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.
-      auto *wrapperDecl =
-          vd->getInterfaceType()->getNominalOrBoundGenericNominal();
-      auto *wrappedValueDecl =
-          wrapperDecl->getPropertyWrapperTypeInfo().valueVar;
-      if (!wrappedValueDecl->getAccessor(AccessorKind::Set))
+      if (!originalProperty->getAccessor(AccessorKind::Set))
         continue;
       // Use the original wrapped property.
       vd = originalProperty;