[AutoDiff] draft of lifting samefile derivative constriant

Author: Marc Rasi

include/swift/AST/DiagnosticsSema.def

@@ -3016,12 +3016,15 @@ NOTE(derivative_attr_result_func_type_mismatch_note,none,
      "%0 does not have expected type %1", (Identifier, Type))
 NOTE(derivative_attr_result_func_original_note,none,
      "%0 defined here", (DeclName))
-ERROR(derivative_attr_not_in_same_file_as_original,none,
-      "derivative not in the same file as the original function", ())
 ERROR(derivative_attr_original_stored_property_unsupported,none,
       "cannot register derivative for stored property %0", (DeclName))
 ERROR(derivative_attr_original_already_has_derivative,none,
       "a derivative already exists for %0", (DeclName))
+ERROR(derivative_attr_visibility_too_broad,none,
+      "derivative function visibility must be at least as restrictive as original function "
+      "visibility", ())
+NOTE(derivative_attr_visibility_too_broad_note,none,
+     "original function defined here", ())
 
 // @transpose
 ERROR(transpose_params_clause_param_not_differentiable,none,

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -96,6 +96,7 @@ class SILDifferentiabilityWitness
   SILDifferentiabilityWitnessKey getKey() const;
   SILModule &getModule() const { return Module; }
   SILLinkage getLinkage() const { return Linkage; }
+  void setLinkage(SILLinkage linkage) { Linkage = linkage; }
   SILFunction *getOriginalFunction() const { return OriginalFunction; }
   const AutoDiffConfig &getConfig() const { return Config; }
   IndexSubset *getParameterIndices() const {
@@ -127,6 +128,7 @@ class SILDifferentiabilityWitness
   bool isDeclaration() const { return IsDeclaration; }
   bool isDefinition() const { return !IsDeclaration; }
   bool isSerialized() const { return IsSerialized; }
+  void setSerialized(bool isSerialized) { IsSerialized = isSerialized; }
   const DeclAttribute *getAttribute() const { return Attribute; }
 
   /// Returns the `SILAutoDiffIndices` corresponding to this config's indices.

lib/SILGen/SILGen.cpp

@@ -781,7 +781,8 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
                "all functions with generic signatures");
         AutoDiffConfig config(diffAttr->getParameterIndices(), resultIndices,
                               diffAttr->getDerivativeGenericSignature());
-        emitDifferentiabilityWitness(AFD, F, config, jvp, vjp, diffAttr);
+        emitDifferentiabilityWitness(AFD, F, config, jvp, vjp, diffAttr,
+                                     F->getLinkage());
       }
       for (auto *derivAttr : Attrs.getAttributes<DerivativeAttr>()) {
         SILFunction *jvp = nullptr;
@@ -801,7 +802,7 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
         AutoDiffConfig config(derivAttr->getParameterIndices(), resultIndices,
                               derivativeGenSig);
         emitDifferentiabilityWitness(origAFD, origFn, config, jvp, vjp,
-                                     derivAttr);
+                                     derivAttr, F->getLinkage());
       }
     };
     if (auto *accessor = dyn_cast<AccessorDecl>(AFD))
@@ -815,7 +816,7 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
 void SILGenModule::emitDifferentiabilityWitness(
     AbstractFunctionDecl *originalAFD, SILFunction *originalFunction,
     const AutoDiffConfig &config, SILFunction *jvp, SILFunction *vjp,
-    const DeclAttribute *attr) {
+    const DeclAttribute *attr, SILLinkage witnessLinkage) {
   assert(isa<DifferentiableAttr>(attr) || isa<DerivativeAttr>(attr));
   auto *origFnType = originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
   auto origSilFnType = originalFunction->getLoweredFunctionType();
@@ -854,11 +855,19 @@ void SILGenModule::emitDifferentiabilityWitness(
   auto *diffWitness = M.lookUpDifferentiabilityWitness(key);
   if (!diffWitness) {
     diffWitness = SILDifferentiabilityWitness::createDefinition(
-        M, originalFunction->getLinkage(), originalFunction,
-        silConfig.parameterIndices, silConfig.resultIndices,
-        config.derivativeGenericSignature, /*jvp*/ nullptr, /*vjp*/ nullptr,
-        /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
-        attr);
+        M, witnessLinkage, originalFunction, silConfig.parameterIndices,
+        silConfig.resultIndices, config.derivativeGenericSignature,
+        /*jvp*/ nullptr, /*vjp*/ nullptr,
+        /*isSerialized*/ hasPublicVisibility(witnessLinkage), attr);
+  }
+
+  // Use the least restrictive declared linkage, so that e.g. a
+  // `@differentiable` on `public` function with `@derivative`s on `internal`
+  // functions results in a public witness. (Sema is responsible for diagnosing
+  // forbidden combinations).
+  if (witnessLinkage < diffWitness->getLinkage()) {
+    diffWitness->setLinkage(witnessLinkage);
+    diffWitness->setSerialized(hasPublicVisibility(witnessLinkage));
   }
 
   // Set derivative function in differentiability witness.

lib/SILGen/SILGen.h

@@ -366,7 +366,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
                                     SILFunction *originalFunction,
                                     const AutoDiffConfig &config,
                                     SILFunction *jvp, SILFunction *vjp,
-                                    const DeclAttribute *diffAttr);
+                                    const DeclAttribute *diffAttr,
+                                    SILLinkage witnessLinkage);
   // SWIFT_ENABLE_TENSORFLOW END
 
   /// Emit the lazy initializer function for a global pattern binding

lib/Sema/TypeCheckAttr.cpp

@@ -3867,12 +3867,15 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
     return;
   }
 
-  // Reject different-file retroactive derivatives.
-  // TODO(TF-136): Lift this restriction now that SIL differentiability witness
-  // infrastructure is ready.
-  if (originalAFD->getParentSourceFile() != derivative->getParentSourceFile()) {
-    diagnoseAndRemoveAttr(attr,
-                          diag::derivative_attr_not_in_same_file_as_original);
+  // Check that derivative visibility is at least as restricted as original
+  // function visibility.
+  if (derivative->getFormalAccessScope() !=
+          originalAFD->getFormalAccessScope() &&
+      !derivative->getFormalAccessScope().isChildOf(
+          originalAFD->getFormalAccessScope())) {
+    diagnoseAndRemoveAttr(attr, diag::derivative_attr_visibility_too_broad);
+    diagnose(originalAFD->getLoc(),
+             diag::derivative_attr_visibility_too_broad_note);
     return;
   }
 

lib/TBDGen/TBDGen.cpp

@@ -67,6 +67,8 @@ void TBDGenVisitor::addSymbol(StringRef name, SymbolKind kind) {
   if (StringSymbols && kind == SymbolKind::GlobalSymbol) {
     auto isNewValue = StringSymbols->insert(mangled).second;
     (void)isNewValue;
+    if (!isNewValue)
+      llvm::dbgs() << mangled << "\n";
     assert(isNewValue && "symbol appears twice");
   }
 }
@@ -236,6 +238,10 @@ void TBDGenVisitor::addDifferentiabilityWitness(
 
 void TBDGenVisitor::addDerivativeConfiguration(AbstractFunctionDecl *original,
                                                AutoDiffConfig config) {
+  auto inserted = AddedDerivatives.insert({original, config});
+  if (!inserted.second)
+    return;
+
   addAutoDiffLinearMapFunction(original, config,
                                AutoDiffLinearMapKind::Differential);
   addAutoDiffLinearMapFunction(original, config,
@@ -315,9 +321,20 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
   }
 
   // SWIFT_ENABLE_TENSORFLOW
-  for (auto derivativeConfig : AFD->getDerivativeFunctionConfigurations()) {
-    addDerivativeConfiguration(AFD, derivativeConfig);
-  }
+  for (const auto *differentiableAttr :
+       AFD->getAttrs().getAttributes<DifferentiableAttr>())
+    addDerivativeConfiguration(
+        AFD,
+        AutoDiffConfig(differentiableAttr->getParameterIndices(),
+                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       differentiableAttr->getDerivativeGenericSignature()));
+  for (const auto *derivativeAttr :
+       AFD->getAttrs().getAttributes<DerivativeAttr>())
+    addDerivativeConfiguration(
+        derivativeAttr->getOriginalFunction(),
+        AutoDiffConfig(derivativeAttr->getParameterIndices(),
+                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       AFD->getGenericSignature()));
 
   visitDefaultArguments(AFD, AFD->getParameters());
 }
@@ -371,6 +388,15 @@ void TBDGenVisitor::visitAbstractStorageDecl(AbstractStorageDecl *ASD) {
   ASD->visitEmittedAccessors([&](AccessorDecl *accessor) {
     visitFuncDecl(accessor);
   });
+
+  // SWIFT_ENABLE_TENSORFLOW
+  for (const auto *differentiableAttr :
+       ASD->getAttrs().getAttributes<DifferentiableAttr>())
+    addDerivativeConfiguration(
+        ASD->getAccessor(AccessorKind::Get),
+        AutoDiffConfig(differentiableAttr->getParameterIndices(),
+                       IndexSubset::get(ASD->getASTContext(), 1, {0}),
+                       differentiableAttr->getDerivativeGenericSignature()));
 }
 
 void TBDGenVisitor::visitVarDecl(VarDecl *VD) {

lib/TBDGen/TBDGenVisitor.h

@@ -55,6 +55,15 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   ModuleDecl *SwiftModule;
   const TBDGenOptions &Opts;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Tracks derivatives that have been added to the TBD.
+  ///
+  /// Different attributes trigger emission of the same derivatives (e.g.
+  /// `@differentiable` and `@derivative(of:)`), so we use this to deduplicate
+  /// the symbols associated with the derivatives in the TBD.
+  llvm::DenseSet<std::pair<AbstractFunctionDecl *, AutoDiffConfig>>
+      AddedDerivatives;
+
 private:
   void addSymbol(StringRef name, llvm::MachO::SymbolKind kind =
                                      llvm::MachO::SymbolKind::GlobalSymbol);