[AutoDiff] lift samefile derivative constriant

include/swift/AST/ASTContext.h

@@ -313,7 +313,7 @@ class ASTContext final {
   // same parameter indices but different derivative generic signatures.
   llvm::DenseMap<
       std::tuple<Decl *, IndexSubset *, AutoDiffDerivativeFunctionKind>,
-      DerivativeAttr *>
+      llvm::SmallPtrSet<DerivativeAttr *, 1>>
       DerivativeAttrs;
 
 private:

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", (DeclNameRef))
 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", ())
 NOTE(derivative_attr_duplicate_note,none,
      "other attribute declared here", ())
 

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -98,6 +98,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 {
@@ -129,6 +130,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

@@ -782,7 +782,8 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
                "all original SIL 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;
@@ -802,7 +803,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))
@@ -816,7 +817,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();
@@ -855,11 +856,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/SILOptimizer/Utils/Differentiation/Common.cpp

@@ -327,10 +327,6 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
   if (existingWitness)
     return existingWitness;
 
-  assert(original->isExternalDeclaration() &&
-         "SILGen should create differentiability witnesses for all function "
-         "definitions with explicit differentiable attributes");
-
   return SILDifferentiabilityWitness::createDeclaration(
       module, SILLinkage::PublicExternal, original,
       minimalConfig->parameterIndices, minimalConfig->resultIndices,