[AutoDiff] diff wit silgen: type simplification, getters (#28398)

Author: marcrasi

lib/SILGen/SILGen.cpp

@@ -767,19 +767,25 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
       constant.kind != SILDeclRef::Kind::DefaultArgGenerator &&
       !constant.isThunk()) {
     auto *AFD = constant.getAbstractFunctionDecl();
-    // Visit all `@differentiable` attributes.
-    for (auto *diffAttr : AFD->getAttrs().getAttributes<DifferentiableAttr>()) {
-      SILFunction *jvp = nullptr;
-      SILFunction *vjp = nullptr;
-      if (auto *jvpDecl = diffAttr->getJVPFunction())
-        jvp = getFunction(SILDeclRef(jvpDecl), NotForDefinition);
-      if (auto *vjpDecl = diffAttr->getVJPFunction())
-        vjp = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
-      auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
-      AutoDiffConfig config(diffAttr->getParameterIndices(), resultIndices,
-                            diffAttr->getDerivativeGenericSignature().getPointer());
-      emitDifferentiabilityWitness(AFD, F, config, jvp, vjp);
-    }
+    auto emitWitnesses = [&](DeclAttributes &Attrs) {
+      for (auto *diffAttr : Attrs.getAttributes<DifferentiableAttr>()) {
+        SILFunction *jvp = nullptr;
+        SILFunction *vjp = nullptr;
+        if (auto *jvpDecl = diffAttr->getJVPFunction())
+          jvp = getFunction(SILDeclRef(jvpDecl), NotForDefinition);
+        if (auto *vjpDecl = diffAttr->getVJPFunction())
+          vjp = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
+        auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
+        AutoDiffConfig config(
+            diffAttr->getParameterIndices(), resultIndices,
+            diffAttr->getDerivativeGenericSignature().getPointer());
+        emitDifferentiabilityWitness(AFD, F, config, jvp, vjp);
+      }
+    };
+    if (auto *accessor = dyn_cast<AccessorDecl>(AFD))
+      if (accessor->isGetter())
+        emitWitnesses(accessor->getStorage()->getAttrs());
+    emitWitnesses(AFD->getAttrs());
   }
   F->verify();
 }
@@ -817,21 +823,6 @@ void SILGenModule::emitDifferentiabilityWitness(
   CanGenericSignature derivativeCanGenSig;
   if (auto derivativeGenSig = config.derivativeGenericSignature)
     derivativeCanGenSig = derivativeGenSig->getCanonicalSignature();
-  // TODO(TF-835): Use simpler derivative generic signature logic below when
-  // type-checking no longer generates implicit `@differentiable` attributes.
-  // See TF-835 for replacement code.
-  if (jvp) {
-    auto jvpCanGenSig = jvp->getLoweredFunctionType()->getSubstGenericSignature();
-    if (!derivativeCanGenSig && jvpCanGenSig)
-      derivativeCanGenSig = jvpCanGenSig;
-    assert(derivativeCanGenSig == jvpCanGenSig);
-  }
-  if (vjp) {
-    auto vjpCanGenSig = vjp->getLoweredFunctionType()->getSubstGenericSignature();
-    if (!derivativeCanGenSig && vjpCanGenSig)
-      derivativeCanGenSig = vjpCanGenSig;
-    assert(derivativeCanGenSig == vjpCanGenSig);
-  }
   // Create new SIL differentiability witness.
   // Witness JVP and VJP are set below.
   // TODO(TF-919): Explore creating serialized differentiability witnesses.

lib/TBDGen/TBDGen.cpp

@@ -195,20 +195,9 @@ void TBDGenVisitor::addDifferentiabilityWitness(
       attr->getParameterIndices(),
       original->getInterfaceType()->castTo<AnyFunctionType>());
 
-  GenericSignature genericSignature = attr->getDerivativeGenericSignature();
-  if (auto *jvpDecl = attr->getJVPFunction()) {
-    assert(!genericSignature ||
-           jvpDecl->getGenericSignature()->isEqual(genericSignature));
-    genericSignature = jvpDecl->getGenericSignature();
-  }
-  if (auto *vjpDecl = attr->getVJPFunction()) {
-    assert(!genericSignature ||
-           vjpDecl->getGenericSignature()->isEqual(genericSignature));
-    genericSignature = vjpDecl->getGenericSignature();
-  }
-
   std::string originalMangledName = SILDeclRef(original).mangle();
-  AutoDiffConfig config{loweredParamIndices, resultIndices, genericSignature};
+  AutoDiffConfig config{loweredParamIndices, resultIndices,
+                        attr->getDerivativeGenericSignature()};
   SILDifferentiabilityWitnessKey key(originalMangledName, config);
 
   Mangle::ASTMangler mangle;

test/AutoDiff/sil_differentiability_witness_silgen.swift

@@ -74,6 +74,10 @@ func generic_vjp<T: Differentiable>(_ x: T, _ y: Float) -> (
 public struct Foo: Differentiable {
   public var x: Float
 
+// CHECK-LABEL: // differentiability witness for Foo.x.getter
+// CHECK-NEXT: sil_differentiability_witness [parameters 0] [results 0] @$s36sil_differentiability_witness_silgen3FooV1xSfvg : $@convention(method) (Foo) -> Float {
+// CHECK-NEXT: }
+
   @differentiable
   public init(_ x: Float) {
     self.x = x