Correctly synthesize semantic member _modify accessor. Support

differentiation of _modify accessor for wrapped values.

Fixes #55084

Author: asl

lib/SILOptimizer/Differentiation/Common.cpp

@@ -61,10 +61,10 @@ bool isSemanticMemberAccessor(SILFunction *original) {
   auto *accessor = dyn_cast<AccessorDecl>(decl);
   if (!accessor)
     return false;
-  // Currently, only getters and setters are supported.
-  // TODO(https://github.com/apple/swift/issues/55084): Support `modify` accessors.
+  // Currently, only getters, setters and _modify accessors are supported.
   if (accessor->getAccessorKind() != AccessorKind::Get &&
-      accessor->getAccessorKind() != AccessorKind::Set)
+      accessor->getAccessorKind() != AccessorKind::Set &&
+      accessor->getAccessorKind() != AccessorKind::Modify)
     return false;
   // Accessor must come from a `var` declaration.
   auto *varDecl = dyn_cast<VarDecl>(accessor->getStorage());

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -905,6 +905,7 @@ class PullbackCloner::Implementation final
   bool runForSemanticMemberAccessor();
   bool runForSemanticMemberGetter();
   bool runForSemanticMemberSetter();
+  bool runForSemanticMemberModify();
 
   /// If original result is non-varied, it will always have a zero derivative.
   /// Skip full pullback generation and simply emit zero derivatives for wrt
@@ -2452,7 +2453,8 @@ bool PullbackCloner::Implementation::run() {
 
   // If the original function is an accessor with special-case pullback
   // generation logic, do special-case generation.
-  if (isSemanticMemberAccessor(&original)) {
+  bool isSemanticMemberAcc = isSemanticMemberAccessor(&original);
+  if (isSemanticMemberAcc) {
     if (runForSemanticMemberAccessor())
       return true;
   }
@@ -2730,7 +2732,8 @@ bool PullbackCloner::Implementation::run() {
 #endif
 
   LLVM_DEBUG(getADDebugStream()
-             << "Generated pullback for " << original.getName() << ":\n"
+             << "Generated " << (isSemanticMemberAcc ? "semantic member accessor" : "normal")
+             << " pullback for " << original.getName() << ":\n"
              << pullback);
   return errorOccurred;
 }
@@ -3205,7 +3208,8 @@ bool PullbackCloner::Implementation::runForSemanticMemberAccessor() {
     return runForSemanticMemberGetter();
   case AccessorKind::Set:
     return runForSemanticMemberSetter();
-  // TODO(https://github.com/apple/swift/issues/55084): Support `modify` accessors.
+  case AccessorKind::Modify:
+    return runForSemanticMemberModify();
   default:
     llvm_unreachable("Unsupported accessor kind; inconsistent with "
                      "`isSemanticMemberAccessor`?");
@@ -3389,6 +3393,83 @@ bool PullbackCloner::Implementation::runForSemanticMemberSetter() {
   return false;
 }
 
+bool PullbackCloner::Implementation::runForSemanticMemberModify() {
+  auto &original = getOriginal();
+  auto &pullback = getPullback();
+  auto pbLoc = getPullback().getLocation();
+
+  auto *accessor = cast<AccessorDecl>(original.getDeclContext()->getAsDecl());
+  assert(accessor->getAccessorKind() == AccessorKind::Modify);
+
+  auto *origEntry = original.getEntryBlock();
+  // We assume that the accessor has a simple 3-BB structure with yield in the entry BB
+  // plus resume and unwind BBs
+  auto *yi = cast<YieldInst>(origEntry->getTerminator());
+  auto *origResumeBB = yi->getResumeBB();
+
+  auto *pbEntry = pullback.getEntryBlock();
+  builder.setCurrentDebugScope(
+      remapScope(origEntry->getScopeOfFirstNonMetaInstruction()));
+  builder.setInsertionPoint(pbEntry);
+
+  // Get _modify accessor argument values.
+  // Accessor type : $(inout Self) -> @yields @inout Argument
+  // Pullback type : $(inout Self', linear map tuple) -> @yields @inout Argument'
+  // Normally pullbacks for semantic member accessors are single BB and
+  // therefore have empty linear map tuple, however, coroutines have a branching
+  // control flow due to possible coroutine abort, so we need to accommodate for
+  // this. We keep branch tracing enums in order not to special case in many
+  // other places. As there is no way to return to coroutine via abort exit, we
+  // essentially "linearize" a coroutine.
+  auto loweredFnTy = original.getLoweredFunctionType();
+  auto pullbackLoweredFnTy = pullback.getLoweredFunctionType();
+
+  assert(loweredFnTy->getNumParameters() == 1 &&
+         loweredFnTy->getNumYields() == 1);
+  assert(pullbackLoweredFnTy->getNumParameters() == 2);
+  assert(pullbackLoweredFnTy->getNumYields() == 1);
+
+  SILValue origSelf = original.getArgumentsWithoutIndirectResults().front();
+
+  SmallVector<SILValue, 8> origFormalResults;
+  collectAllFormalResultsInTypeOrder(original, origFormalResults);
+
+  assert(getConfig().resultIndices->getNumIndices() == 2 &&
+         "Modify accessor should have two semantic results");
+
+  auto origYield = origFormalResults[*std::next(getConfig().resultIndices->begin())];
+
+  // Look up the corresponding field in the tangent space.
+  auto *origField = cast<VarDecl>(accessor->getStorage());
+  auto baseType = remapType(origSelf->getType()).getASTType();
+  auto *tanField = getTangentStoredProperty(getContext(), origField, baseType,
+                                            pbLoc, getInvoker());
+  if (!tanField) {
+    errorOccurred = true;
+    return true;
+  }
+
+  auto adjSelf = getAdjointBuffer(origResumeBB, origSelf);
+  auto *adjSelfElt = builder.createStructElementAddr(pbLoc, adjSelf, tanField);
+  // Modify accessors have inout yields and therefore should yield addresses.
+  assert(getTangentValueCategory(origYield) == SILValueCategory::Address &&
+         "Modify accessors should yield indirect");
+
+  // Yield the adjoint buffer and do everything else in the resume
+  // destination. Unwind destination is unreachable as the coroutine can never
+  // be aborted.
+  auto *unwindBB = getPullback().createBasicBlock();
+  auto *resumeBB = getPullbackBlock(origEntry);
+  builder.createYield(yi->getLoc(), {adjSelfElt}, resumeBB, unwindBB);
+  builder.setInsertionPoint(unwindBB);
+  builder.createUnreachable(SILLocation::invalid());
+
+  builder.setInsertionPoint(resumeBB);
+  addToAdjointBuffer(origEntry, origSelf, adjSelf, pbLoc);
+
+  return false;
+}
+
 //--------------------------------------------------------------------------//
 // Adjoint buffer mapping
 //--------------------------------------------------------------------------//

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -460,6 +460,16 @@ class VJPCloner::Implementation final
       TypeSubstCloner::visitEndApplyInst(eai);
       return;
     }
+    // If the original function is a semantic member accessor, do standard
+    // cloning. Semantic member accessors have special pullback generation
+    // logic, so all `end_apply` instructions can be directly cloned to the VJP.
+    if (isSemanticMemberAccessor(original)) {
+      LLVM_DEBUG(getADDebugStream()
+                 << "Cloning `end_apply` in semantic member accessor:\n"
+                 << *eai << '\n');
+      TypeSubstCloner::visitEndApplyInst(eai);
+      return;
+    }
 
     Builder.setCurrentDebugScope(getOpScope(eai->getDebugScope()));
     auto loc = eai->getLoc();
@@ -607,6 +617,16 @@ class VJPCloner::Implementation final
       TypeSubstCloner::visitBeginApplyInst(bai);
       return;
     }
+    // If the original function is a semantic member accessor, do standard
+    // cloning. Semantic member accessors have special pullback generation
+    // logic, so all `begin_apply` instructions can be directly cloned to the VJP.
+    if (isSemanticMemberAccessor(original)) {
+      LLVM_DEBUG(getADDebugStream()
+                 << "Cloning `begin_apply` in semantic member accessor:\n"
+                 << *bai << '\n');
+      TypeSubstCloner::visitBeginApplyInst(bai);
+      return;
+    }
 
     Builder.setCurrentDebugScope(getOpScope(bai->getDebugScope()));
     auto loc = bai->getLoc();

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -680,9 +680,7 @@ extension DifferentiableWrapper: Differentiable where Value: Differentiable {}
 // accesses.
 
 struct Struct: Differentiable {
-  // expected-error @+4 {{expression is not differentiable}}
-  // expected-error @+3 {{expression is not differentiable}}
-  // expected-note @+2 {{cannot differentiate access to property 'Struct._x' because 'Struct.TangentVector' does not have a stored property named '_x'}}
+  // expected-error @+2 {{expression is not differentiable}}
   // expected-note @+1 {{cannot differentiate access to property 'Struct._x' because 'Struct.TangentVector' does not have a stored property named '_x'}}  
   @DifferentiableWrapper @DifferentiableWrapper var x: Float = 10
 

test/AutoDiff/validation-test/property_wrappers.swift

@@ -19,6 +19,7 @@ struct Wrapper<Value> {
   var wrappedValue: Value { // computed property
     get { value }
     set { value = newValue }
+    _modify { yield &value }
   }
 
   init(wrappedValue: Value) {
@@ -46,16 +47,13 @@ PropertyWrapperTests.test("SimpleStruct") {
   expectEqual((.init(x: 60, y: 0, z: 20), 300),
               gradient(at: Struct(), 2, of: setter))
 
-  // TODO: Support `modify` accessors (https://github.com/apple/swift/issues/55084).
-  /*
   func modify(_ s: Struct, _ x: Tracked<Float>) -> Tracked<Float> {
     var s = s
     s.x *= x * s.z
     return s.x
   }
   expectEqual((.init(x: 60, y: 0, z: 20), 300),
               gradient(at: Struct(), 2, of: modify))
-  */
 }
 
 struct GenericStruct<T> {
@@ -86,16 +84,13 @@ PropertyWrapperTests.test("GenericStruct") {
   expectEqual((.init(x: 60, y: 0, z: 20), 300),
               gradient(at: GenericStruct<Tracked<Float>>(y: 20), 2, of: setter))
 
-  // TODO: Support `modify` accessors (https://github.com/apple/swift/issues/55084).
-  /*
   func modify<T>(_ s: GenericStruct<T>, _ x: Tracked<Float>) -> Tracked<Float> {
     var s = s
     s.x *= x * s.z
     return s.x
   }
   expectEqual((.init(x: 60, y: 0, z: 20), 300),
               gradient(at: GenericStruct<Tracked<Float>>(y: 1), 2, of: modify))
-  */
 }
 
 // TF-1149: Test class with loadable type but address-only `TangentVector` type.
@@ -131,16 +126,18 @@ PropertyWrapperTests.test("SimpleClass") {
               gradient(at: Class(), 2, of: setter))
   */
 
-  // TODO: Support `modify` accessors (https://github.com/apple/swift/issues/55084).
-  /*
+  // FIXME(TF-1175): Same issue as above
   func modify(_ c: Class, _ x: Tracked<Float>) -> Tracked<Float> {
     var c = c
     c.x *= x * c.z
     return c.x
   }
+  /*
   expectEqual((.init(x: 60, y: 0, z: 20), 300),
               gradient(at: Class(), 2, of: modify))
   */
+  expectEqual((.init(x: 1, y: 0, z: 0), 0),
+              gradient(at: Class(), 2, of: modify))
 }
 
 // From: https://github.com/apple/swift-evolution/blob/master/proposals/0258-property-wrappers.md#proposed-solution
@@ -157,12 +154,13 @@ enum Lazy<Value> {
 
   var wrappedValue: Value {
     // TODO(TF-1250): Replace with actual mutating getter implementation.
-    // Requires differentiation to support functions with multiple results.
-    get {
+    // Requires support for mutating semantic member accessor
+    /* mutating */ get {
       switch self {
       case .uninitialized(let initializer):
         let value = initializer()
         // NOTE: Actual implementation assigns to `self` here.
+        // self = .initialized(value)
         return value
       case .initialized(let value):
         return value