Don't allow marking stored props as differentiable

Author: bartchr808

include/swift/AST/DiagnosticsSema.def

@@ -2733,6 +2733,8 @@ ERROR(differentiable_attr_unsupported_req_kind,none,
       "layout requirement are not supported by '@differentiable' attribute", ())
 ERROR(differentiable_attr_class_unsupported,none,
       "class members cannot be marked with '@differentiable'", ())
+ERROR(differentiable_attr_stored_prop_unsupported,none,
+"Stored properties cannot be marked with '@differentiable'", ())
 NOTE(protocol_witness_missing_specific_differentiable_attr,none,
      "candidate is missing attribute '%0'", (StringRef))
 

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -3234,27 +3234,22 @@ class VJPEmitter final
     }
     // 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 (structDecl->getEffectiveAccess() <= AccessLevel::Internal ||
+    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;
-    // Find the corresponding getter and its VJP.
-    auto *getterDecl = sei->getField()->getGetter();
-    assert(getterDecl);
-    auto *getterFn = getModule().lookUpFunction(
-        SILDeclRef(getterDecl, SILDeclRef::Kind::Func));
-    if (!getterFn) {
-      context.emitNondifferentiabilityError(
-          sei, invoker, diag::autodiff_property_not_differentiable);
-      errorOccurred = true;
-      return;
-    }
     SILAutoDiffIndices indices(/*source*/ 0,
         AutoDiffIndexSubset::getDefault(getASTContext(), 1, true));
     auto *attr = context.lookUpDifferentiableAttr(getterFn, indices);
@@ -3304,27 +3299,22 @@ class VJPEmitter final
     }
     // 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 (structDecl->getEffectiveAccess() <= AccessLevel::Internal ||
+    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;
-    // Find the corresponding getter and its VJP.
-    auto *getterDecl = seai->getField()->getGetter();
-    assert(getterDecl);
-    auto *getterFn = getModule().lookUpFunction(
-        SILDeclRef(getterDecl, SILDeclRef::Kind::Func));
-    if (!getterFn) {
-      context.emitNondifferentiabilityError(
-          seai, invoker, diag::autodiff_property_not_differentiable);
-      errorOccurred = true;
-      return;
-    }
     SILAutoDiffIndices indices(/*source*/ 0,
         AutoDiffIndexSubset::getDefault(getASTContext(), 1, true));
     auto *attr = context.lookUpDifferentiableAttr(getterFn, indices);

lib/Sema/TypeCheckAttr.cpp

@@ -2887,6 +2887,9 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
 
   AbstractFunctionDecl *original = dyn_cast<AbstractFunctionDecl>(D);
   if (auto *asd = dyn_cast<AbstractStorageDecl>(D)) {
+    if (asd->getImplInfo().isSimpleStored()) {
+      diagnoseAndRemoveAttr(attr, diag::differentiable_attr_stored_prop_unsupported);
+    }
     // When used directly on a storage decl (stored/computed property or
     // subscript), the getter is currently inferred to be `@differentiable`.
     // TODO(TF-129): Infer setter to also be `@differentiable` after

test/AutoDiff/differentiating_attr_type_checking.swift

@@ -295,3 +295,23 @@ func jvpConsistent(_ x: Float) -> (value: Float, differential: (Float) -> Float)
 func vjpConsistent(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
   return (x, { $0 })
 }
+
+// Test usage of `@differentiable` on a stored property
+struct PropertyDiff : Differentiable & AdditiveArithmetic {
+	// expected-error @+1 {{Stored properties cannot be marked with '@differentiable'}}
+    @differentiable(vjp: vjpPropertyA)
+    var a: Float = 1
+    typealias TangentVector = PropertyDiff
+    typealias AllDifferentiableVariables = PropertyDiff
+    func vjpPropertyA() -> (Float, (Float) -> PropertyDiff) {
+        (.zero, { _ in .zero })
+    }
+}
+
+@differentiable
+func f(_ x: PropertyDiff) -> Float {
+    return x.a
+}
+
+let a = gradient(at: PropertyDiff(), in: f)
+print(a)