Remove getter differentiation pass, fix tests.

Author: bartchr808

include/swift/AST/DiagnosticsSema.def

@@ -2734,7 +2734,7 @@ ERROR(differentiable_attr_unsupported_req_kind,none,
 ERROR(differentiable_attr_class_unsupported,none,
       "class members cannot be marked with '@differentiable'", ())
 ERROR(differentiable_attr_stored_property_variable_unsupported,none,
-      "stored properties/variables cannot be marked with '@differentiable' with a custom VJP/JVP", ())
+      "'jvp:' or 'vjp:' cannot be specified for stored properties", ())
 NOTE(protocol_witness_missing_specific_differentiable_attr,none,
      "candidate is missing attribute '%0'", (StringRef))
 

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -737,11 +737,7 @@ enum class StructExtractDifferentiationStrategy {
   // that is zero except along the direction of the corresponding field.
   //
   // Fields correspond by matching name.
-  Fieldwise,
-
-  // Differentiate the `struct_extract` by looking up the corresponding getter
-  // and using its VJP.
-  Getter
+  Fieldwise
 };
 
 static inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
@@ -3232,59 +3228,10 @@ class VJPEmitter final
       SILClonerWithScopes::visitStructExtractInst(sei);
       return;
     }
-    // This instruction is active. Determine the appropriate differentiation
-    // strategy, and use it.
-    // Find the corresponding getter.
-    auto *getterDecl = sei->getField()->getGetter();
-    assert(getterDecl);
-    auto *getterFn = getModule().lookUpFunction(
-        SILDeclRef(getterDecl, SILDeclRef::Kind::Func));
-    auto *structDecl = sei->getStructDecl();
-    if (!getterFn ||
-        structDecl->getAttrs().hasAttribute<FieldwiseDifferentiableAttr>()) {
-      strategies[sei] = StructExtractDifferentiationStrategy::Fieldwise;
-      SILClonerWithScopes::visitStructExtractInst(sei);
-      return;
-    }
-    // The FieldwiseProductSpace strategy is not appropriate, so use the Getter
-    // strategy.
-    assert(getterFn);
-    strategies[sei] = StructExtractDifferentiationStrategy::Getter;
-    SILAutoDiffIndices indices(/*source*/ 0,
-        AutoDiffIndexSubset::getDefault(getASTContext(), 1, true));
-    auto *attr = context.lookUpDifferentiableAttr(getterFn, indices);
-    if (!attr) {
-      context.emitNondifferentiabilityError(
-          sei, invoker, diag::autodiff_property_not_differentiable);
-      errorOccurred = true;
-      return;
-    }
-    // Reference and apply the VJP.
-    auto loc = sei->getLoc();
-    auto *getterVJP = getAssociatedFunction(
-        context, getterFn, attr, AutoDiffAssociatedFunctionKind::VJP,
-        attr->getVJPName());
-    assert(getterVJP && "Expected to find getter VJP");
-    auto *getterVJPRef = getBuilder().createFunctionRef(loc, getterVJP);
-    auto *getterVJPApply = getBuilder().createApply(
-        loc, getterVJPRef,
-        getOpSubstitutionMap(getterVJP->getForwardingSubstitutionMap()),
-        /*args*/ {getOpValue(sei->getOperand())}, /*isNonThrowing*/ false);
-    // Extract direct results from `getterVJPApply`.
-    SmallVector<SILValue, 8> vjpDirectResults;
-    extractAllElements(getterVJPApply, getBuilder(), vjpDirectResults);
-    // Map original result.
-    auto originalDirectResults =
-        ArrayRef<SILValue>(vjpDirectResults).drop_back(1);
-    auto originalDirectResult = joinElements(originalDirectResults,
-                                             getBuilder(),
-                                             getterVJPApply->getLoc());
-    mapValue(sei, originalDirectResult);
-    // Checkpoint the pullback.
-    auto pullback = vjpDirectResults.back();
-    // TODO: Check whether it's necessary to reabstract getter pullbacks.
-    pullbackInfo.addPullbackDecl(sei, getOpType(pullback->getType()));
-    pullbackValues[sei->getParent()].push_back(pullback);
+    // This instruction is active. Use the field wise differentiation strategy
+    // to differentiate the struct extract instruction.
+    strategies[sei] = StructExtractDifferentiationStrategy::Fieldwise;
+    SILClonerWithScopes::visitStructExtractInst(sei);
   }
 
   void visitStructElementAddrInst(StructElementAddrInst *seai) {
@@ -3297,78 +3244,10 @@ class VJPEmitter final
       SILClonerWithScopes::visitStructElementAddrInst(seai);
       return;
     }
-    // This instruction is active. Determine the appropriate differentiation
-    // strategy, and use it.
-    // Find the corresponding getter.
-    auto *getterDecl = seai->getField()->getGetter();
-    assert(getterDecl);
-    auto *getterFn = getModule().lookUpFunction(
-        SILDeclRef(getterDecl, SILDeclRef::Kind::Func));
-    auto *structDecl = seai->getStructDecl();
-    if (!getterFn ||
-        structDecl->getAttrs().hasAttribute<FieldwiseDifferentiableAttr>()) {
-      strategies[seai] = StructExtractDifferentiationStrategy::Fieldwise;
-      SILClonerWithScopes::visitStructElementAddrInst(seai);
-      return;
-    }
-    // The FieldwiseProductSpace strategy is not appropriate, so use the Getter
-    // strategy.
-    assert(getterFn);
-    strategies[seai] = StructExtractDifferentiationStrategy::Getter;
-    SILAutoDiffIndices indices(/*source*/ 0,
-        AutoDiffIndexSubset::getDefault(getASTContext(), 1, true));
-    auto *attr = context.lookUpDifferentiableAttr(getterFn, indices);
-    if (!attr) {
-      context.emitNondifferentiabilityError(
-          seai, invoker, diag::autodiff_property_not_differentiable);
-      errorOccurred = true;
-      return;
-    }
-    // Set generic context scope before getting VJP function type.
-    auto vjpGenSig = SubsMap.getGenericSignature()
-        ? SubsMap.getGenericSignature()->getCanonicalSignature()
-        : nullptr;
-    Lowering::GenericContextScope genericContextScope(
-        context.getTypeConverter(), vjpGenSig);
-    // Reference the getter VJP.
-    auto loc = seai->getLoc();
-    auto *getterVJP = getModule().lookUpFunction(attr->getVJPName());
-    assert(getterVJP && "Expected to find getter VJP");
-    auto vjpFnTy = getterVJP->getLoweredFunctionType();
-    auto *getterVJPRef = getBuilder().createFunctionRef(loc, getterVJP);
-    // Store getter VJP arguments and indirect result buffers.
-    SmallVector<SILValue, 8> vjpArgs;
-    SmallVector<AllocStackInst *, 8> vjpIndirectResults;
-    for (auto indRes : vjpFnTy->getIndirectFormalResults()) {
-      auto *alloc = getBuilder().createAllocStack(
-          loc, getOpType(indRes.getSILStorageType()));
-      vjpArgs.push_back(alloc);
-      vjpIndirectResults.push_back(alloc);
-    }
-    vjpArgs.push_back(getOpValue(seai->getOperand()));
-    // Apply the getter VJP.
-    auto *getterVJPApply = getBuilder().createApply(
-        loc, getterVJPRef,
-        getOpSubstitutionMap(getterVJP->getForwardingSubstitutionMap()),
-        vjpArgs, /*isNonThrowing*/ false);
-    // Collect all results from `getterVJPApply` in type-defined order.
-    SmallVector<SILValue, 8> vjpDirectResults;
-    extractAllElements(getterVJPApply, getBuilder(), vjpDirectResults);
-    SmallVector<SILValue, 8> allResults;
-    collectAllActualResultsInTypeOrder(
-        getterVJPApply, vjpDirectResults,
-        getterVJPApply->getIndirectSILResults(), allResults);
-    // Deallocate VJP indirect results.
-    for (auto alloc : vjpIndirectResults)
-      getBuilder().createDeallocStack(loc, alloc);
-    auto originalDirectResult = allResults[indices.source];
-    // Map original result.
-    mapValue(seai, originalDirectResult);
-    // Checkpoint the pullback.
-    SILValue pullback = vjpDirectResults.back();
-    // TODO: Check whether it's necessary to reabstract getter pullbacks.
-    pullbackInfo.addPullbackDecl(seai, getOpType(pullback->getType()));
-    pullbackValues[seai->getParent()].push_back(pullback);
+    // This instruction is active. Use the field wise differentiation strategy
+    // to differentiate the struct extract instruction.
+    strategies[seai] = StructExtractDifferentiationStrategy::Fieldwise;
+    SILClonerWithScopes::visitStructElementAddrInst(seai);
   }
 
   // If an `apply` has active results or active inout parameters, replace it
@@ -4839,29 +4718,6 @@ class AdjointEmitter final : public SILInstructionVisitor<AdjointEmitter> {
       }
       return;
     }
-    case StructExtractDifferentiationStrategy::Getter: {
-      // Get the pullback.
-      auto *pullbackField = getPullbackInfo().lookUpPullbackDecl(sei);
-      assert(pullbackField);
-      auto pullback = builder.createStructExtract(
-          loc, getAdjointBlockPullbackStructArgument(sei->getParent()),
-          pullbackField);
-
-      // Construct the pullback arguments.
-      auto av = takeAdjointValue(sei);
-      auto vector = materializeAdjointDirect(std::move(av), loc);
-
-      // Call the pullback.
-      auto *pullbackCall = builder.createApply(
-          loc, pullback, SubstitutionMap(), {vector}, /*isNonThrowing*/ false);
-      assert(!pullbackCall->hasIndirectResults());
-
-      // Accumulate adjoint for the `struct_extract` operand.
-      addAdjointValue(sei->getOperand(),
-          makeConcreteAdjointValue(
-              ValueWithCleanup(pullbackCall, vector.getCleanup())));
-      break;
-    }
     }
   }
 

lib/Sema/TypeCheckAttr.cpp

@@ -2891,7 +2891,7 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
         (attr->getJVP() || attr->getVJP())) {
       diagnoseAndRemoveAttr(attr,
           diag::differentiable_attr_stored_property_variable_unsupported);
-	  return;
+      return;
     }
     // When used directly on a storage decl (stored/computed property or
     // subscript), the getter is currently inferred to be `@differentiable`.

test/AutoDiff/autodiff_diagnostics.swift

@@ -38,8 +38,6 @@ extension S : Differentiable, VectorNumeric {
   typealias TangentVector = S
 }
 
-// expected-error @+2 {{function is not differentiable}}
-// expected-note @+1 {{property is not differentiable}}
 _ = gradient(at: S(p: 0)) { s in 2 * s.p }
 
 struct NoDerivativeProperty : Differentiable {

test/AutoDiff/differentiable_attr_silgen.swift

@@ -76,43 +76,6 @@ public func dhasvjp(_ x: Float, _ y: Float) -> (Float, (Float) -> (Float, Float)
 
 // CHECK-LABEL: sil [ossa] @dhasvjp
 
-//===----------------------------------------------------------------------===//
-// Stored property
-//===----------------------------------------------------------------------===//
-
-struct DiffStoredProp {
-  @differentiable(wrt: (self), jvp: storedPropJVP, vjp: storedPropVJP)
-  let storedProp: Float
-
-  @_silgen_name("storedPropJVP")
-  func storedPropJVP() -> (Float, (DiffStoredProp) -> Float) {
-    fatalError("unimplemented")
-  }
-
-  @_silgen_name("storedPropVJP")
-  func storedPropVJP() -> (Float, (Float) -> DiffStoredProp) {
-    fatalError("unimplemented")
-  }
-}
-
-extension DiffStoredProp : VectorNumeric {
-  static var zero: DiffStoredProp { fatalError("unimplemented") }
-  static func + (lhs: DiffStoredProp, rhs: DiffStoredProp) -> DiffStoredProp {
-    fatalError("unimplemented")
-  }
-  static func - (lhs: DiffStoredProp, rhs: DiffStoredProp) -> DiffStoredProp {
-    fatalError("unimplemented")
-  }
-  typealias Scalar = Float
-  static func * (lhs: Float, rhs: DiffStoredProp) -> DiffStoredProp {
-    fatalError("unimplemented")
-  }
-}
-
-extension DiffStoredProp : Differentiable {
-  typealias TangentVector = DiffStoredProp
-}
-
 //===----------------------------------------------------------------------===//
 // Computed property
 //===----------------------------------------------------------------------===//

test/AutoDiff/differentiable_attr_type_checking.swift

@@ -1,7 +1,10 @@
 // RUN: %target-swift-frontend -typecheck -verify %s
 
 @differentiable // expected-error {{'@differentiable' attribute cannot be applied to this declaration}}
-let global: Float = 1
+let globalConst: Float = 1
+
+@differentiable // expected-error {{'@differentiable' attribute cannot be applied to this declaration}}
+var globalVar: Float = 1
 
 func testLocalVariables() {
   // expected-error @+1 {{'_' has no parameters to differentiate with respect to}}
@@ -225,25 +228,18 @@ class Foo {
 }
 
 struct JVPStruct {
+  @differentiable
   let p: Float
 
-  @differentiable(wrt: (self), jvp: storedPropJVP)
-  let storedImmutableOk: Float
-
-  // expected-error @+1 {{'storedPropJVP' does not have expected type '(JVPStruct) -> () -> (Double, (JVPStruct.TangentVector) -> Double.TangentVector)' (aka '(JVPStruct) -> () -> (Double, (JVPStruct) -> Double)'}}
-  @differentiable(wrt: (self), jvp: storedPropJVP)
-  let storedImmutableWrongType: Double
-
-  @differentiable(wrt: (self), jvp: storedPropJVP)
-  var storedMutableOk: Float
-
-  // expected-error @+1 {{'storedPropJVP' does not have expected type '(JVPStruct) -> () -> (Double, (JVPStruct.TangentVector) -> Double.TangentVector)' (aka '(JVPStruct) -> () -> (Double, (JVPStruct) -> Double)'}}
-  @differentiable(wrt: (self), jvp: storedPropJVP)
-  var storedMutableWrongType: Double
+  // expected-error @+1 {{'funcJVP' does not have expected type '(JVPStruct) -> () -> (Double, (JVPStruct.TangentVector) -> Double.TangentVector)' (aka '(JVPStruct) -> () -> (Double, (JVPStruct) -> Double)'}}
+  @differentiable(wrt: (self), jvp: funcJVP)
+  func funcWrongType() -> Double {
+    fatalError("unimplemented")
+  }
 }
 
 extension JVPStruct {
-  func storedPropJVP() -> (Float, (JVPStruct) -> Float) {
+  func funcJVP() -> (Float, (JVPStruct) -> Float) {
     fatalError("unimplemented")
   }
 }
@@ -383,23 +379,15 @@ func vjpNonDiffResult2(x: Float) -> (Float, Int) {
 struct VJPStruct {
   let p: Float
 
-  @differentiable(vjp: storedPropVJP)
-  let storedImmutableOk: Float
-
-  // expected-error @+1 {{'storedPropVJP' does not have expected type '(VJPStruct) -> () -> (Double, (Double.TangentVector) -> VJPStruct.TangentVector)' (aka '(VJPStruct) -> () -> (Double, (Double) -> VJPStruct)'}}
-  @differentiable(vjp: storedPropVJP)
-  let storedImmutableWrongType: Double
-
-  @differentiable(vjp: storedPropVJP)
-  var storedMutableOk: Float
-
-  // expected-error @+1 {{'storedPropVJP' does not have expected type '(VJPStruct) -> () -> (Double, (Double.TangentVector) -> VJPStruct.TangentVector)' (aka '(VJPStruct) -> () -> (Double, (Double) -> VJPStruct)'}}
-  @differentiable(vjp: storedPropVJP)
-  var storedMutableWrongType: Double
+  // expected-error @+1 {{'funcVJP' does not have expected type '(VJPStruct) -> () -> (Double, (Double.TangentVector) -> VJPStruct.TangentVector)' (aka '(VJPStruct) -> () -> (Double, (Double) -> VJPStruct)'}}
+  @differentiable(vjp: funcVJP)
+  func funcWrongType() -> Double {
+    fatalError("unimplemented")
+  }
 }
 
 extension VJPStruct {
-  func storedPropVJP() -> (Float, (Float) -> VJPStruct) {
+  func funcVJP() -> (Float, (Float) -> VJPStruct) {
     fatalError("unimplemented")
   }
 }

test/AutoDiff/differentiating_attr_type_checking.swift

@@ -298,7 +298,7 @@ func vjpConsistent(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
 
 // Test usage of `@differentiable` on a stored property
 struct PropertyDiff : Differentiable & AdditiveArithmetic {
-    // expected-error @+1 {{stored properties cannot be marked with '@differentiable'}}
+    // expected-error @+1 {{'jvp:' or 'vjp:' cannot be specified for stored properties}}
     @differentiable(vjp: vjpPropertyA)
     var a: Float = 1
     typealias TangentVector = PropertyDiff

test/AutoDiff/e2e_differentiable_property.swift

@@ -30,10 +30,9 @@ struct Space {
   }
 
   private let storedX: Float
-
-  /// `y` is a stored property with a custom vjp for its getter.
-  @differentiable(vjp: vjpY)
-  let y: Float
+  
+  @differentiable
+  var y: Float
 
   func vjpY() -> (Float, (Float) -> TangentSpace) {
     return (y, { v in TangentSpace(dx: 0, dy: v) })
@@ -70,7 +69,7 @@ E2EDifferentiablePropertyTests.test("stored property") {
 
 struct GenericMemberWrapper<T : Differentiable> : Differentiable {
   // Stored property.
-  @differentiable(vjp: vjpX)
+  @differentiable
   var x: T
 
   func vjpX() -> (T, (T.TangentVector) -> GenericMemberWrapper.TangentVector) {