[AutoDiff] improve symbol linkage (#28582)

Changed the linkage and serialization of witnesses, derivative functions, and closures to match the design.
Made deserialized witnesses from other modules not get emitted into the current module IR, by setting their linkage to public_external, and teaching IRGen not to emit public_external witnesses. (Fixes TF-1025).
Added diagnostic for calling private derivatives from serialized functions (Fixes TF-690).

Author: marcrasi

include/swift/AST/DiagnosticsSIL.def

@@ -472,6 +472,12 @@ NOTE(autodiff_expression_not_differentiable_note,none,
 NOTE(autodiff_external_nondifferentiable_function,none,
      "cannot differentiate functions that have not been marked "
      "'@differentiable' and that are defined in other files", ())
+NOTE(autodiff_private_derivative_from_fragile,none,
+     "differentiated functions in "
+     "%select{'@inlinable' functions|default arguments}0 must be marked "
+     "'@differentiable' or have a public '@derivative'"
+     "%select{|; this is not possible with a closure, make a top-level "
+     "function instead}1", (unsigned, bool))
 NOTE(autodiff_nondifferentiable_argument,none,
      "cannot differentiate through a non-differentiable argument; do you want "
      "to use 'withoutDerivative(at:)'?", ())

lib/IRGen/GenDiffWitness.cpp

@@ -34,6 +34,11 @@ void IRGenModule::emitSILDifferentiabilityWitness(
   if (dw->isDeclaration())
     return;
 
+  // Don't emit public_external witnesses.
+  if (hasPublicVisibility(dw->getLinkage()) &&
+      isAvailableExternally(dw->getLinkage()))
+    return;
+
   ConstantInitBuilder builder(*this);
   auto diffWitnessContents = builder.beginStruct();
 

lib/SILGen/SILGen.cpp

@@ -823,7 +823,8 @@ void SILGenModule::emitDifferentiabilityWitness(
   auto *diffWitness = SILDifferentiabilityWitness::createDefinition(
       M, originalFunction->getLinkage(), originalFunction, loweredParamIndices,
       config.resultIndices, config.derivativeGenericSignature,
-      /*jvp*/ nullptr, /*vjp*/ nullptr, originalFunction->isSerialized(),
+      /*jvp*/ nullptr, /*vjp*/ nullptr,
+      /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
       diffAttr);
 
   // Set derivative function in differentiability witness.

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -209,15 +209,16 @@ class DifferentiationTransformer {
 
   ADContext &getContext() { return context; }
 
-  /// Canonicalize the given witness, filling in JVP/VJPs if missing.
+  /// Canonicalize the given witness, filling in derivative functions if
+  /// missing.
   ///
-  /// \param explicitDifferentiable specifies whether the witness comes from an
-  ///        explicit `@differentiable` or `@derivative` attribute in the AST.
-  ///        If it does, we emit JVP/VJPs with the same linkage as the original
-  ///        so that they are linkable from other modules.
+  /// Generated derivative functions have the same linkage as the witness.
+  ///
+  /// \param serializeFunctions specifies whether generated functions should be
+  ///        serialized.
   bool canonicalizeDifferentiabilityWitness(
       SILFunction *original, SILDifferentiabilityWitness *witness,
-      DifferentiationInvoker invoker, bool explicitDifferentiable);
+      DifferentiationInvoker invoker, IsSerialized_t serializeFunctions);
 
   /// Process the given `differentiable_function` instruction, filling in
   /// missing derivative functions if necessary.
@@ -692,17 +693,31 @@ emitDerivativeFunctionReference(
               originalFn->getLoweredFunctionType(), desiredParameterIndices,
               contextualDerivativeGenSig);
       minimalWitness = SILDifferentiabilityWitness::createDefinition(
-          context.getModule(),
-          originalFn->isSerialized() ? SILLinkage::Shared : SILLinkage::Hidden,
-          originalFn, desiredParameterIndices, desiredResultIndices,
+          context.getModule(), SILLinkage::Private, originalFn,
+          desiredParameterIndices, desiredResultIndices,
           derivativeConstrainedGenSig, /*jvp*/ nullptr,
-          /*vjp*/ nullptr, originalFn->isSerialized());
+          /*vjp*/ nullptr, /*isSerialized*/ false);
       if (transformer.canonicalizeDifferentiabilityWitness(
-              originalFn, minimalWitness, invoker,
-              /*explicitDifferentiable*/ false))
+              originalFn, minimalWitness, invoker, IsNotSerialized))
         return None;
     }
     assert(minimalWitness);
+    if (original->getFunction()->isSerialized() &&
+        !hasPublicVisibility(minimalWitness->getLinkage())) {
+      enum { Inlinable = 0, DefaultArgument = 1 };
+      unsigned fragileKind = Inlinable;
+      // FIXME: This is not a very robust way of determining if the function is
+      // a default argument. Also, we have not exhaustively listed all the kinds
+      // of fragility.
+      if (original->getFunction()->getLinkage() == SILLinkage::PublicNonABI)
+        fragileKind = DefaultArgument;
+      context.emitNondifferentiabilityError(
+          original, invoker, diag::autodiff_private_derivative_from_fragile,
+          fragileKind,
+          llvm::isa_and_nonnull<AbstractClosureExpr>(
+              originalFRI->getLoc().getAsASTNode<Expr>()));
+      return None;
+    }
     // TODO(TF-482): Move generic requirement checking logic to
     // `getExactDifferentiabilityWitness` &
     // `getOrCreateMinimalASTDifferentiabilityWitness`.
@@ -1503,12 +1518,11 @@ class VJPEmitter final
         original->getASTContext());
 
     SILOptFunctionBuilder fb(context.getTransform());
-    // The generated pullback linkage is set to Hidden because generated
-    // pullbacks are never called cross-module.
-    auto linkage = SILLinkage::Hidden;
+    auto linkage =
+        vjp->isSerialized() ? SILLinkage::Public : SILLinkage::Private;
     auto *pullback = fb.createFunction(
         linkage, pbName, pbType, pbGenericEnv, original->getLocation(),
-        original->isBare(), IsNotTransparent, original->isSerialized(),
+        original->isBare(), IsNotTransparent, vjp->isSerialized(),
         original->isDynamicallyReplaceable());
     pullback->setDebugScope(new (module)
                                 SILDebugScope(original->getLocation(),
@@ -3114,18 +3128,20 @@ class JVPEmitter final
                                      witness->getSILAutoDiffIndices(), jvp)),
         differentialInfo(context, AutoDiffLinearMapKind::Differential, original,
                          jvp, witness->getSILAutoDiffIndices(), activityInfo),
-        differentialBuilder(SILBuilder(*createEmptyDifferential(
-            context, original, witness, &differentialInfo))),
+        differentialBuilder(SILBuilder(
+            *createEmptyDifferential(context, witness, &differentialInfo))),
         diffLocalAllocBuilder(getDifferential()) {
     // Create empty differential function.
     context.recordGeneratedFunction(&getDifferential());
   }
 
   static SILFunction *
-  createEmptyDifferential(ADContext &context, SILFunction *original,
+  createEmptyDifferential(ADContext &context,
                           SILDifferentiabilityWitness *witness,
                           LinearMapInfo *linearMapInfo) {
     auto &module = context.getModule();
+    auto *original = witness->getOriginalFunction();
+    auto *jvp = witness->getJVP();
     auto origTy = original->getLoweredFunctionType();
     auto lookupConformance = LookUpConformanceInModule(module.getSwiftModule());
 
@@ -3186,12 +3202,11 @@ class JVPEmitter final
         original->getASTContext());
 
     SILOptFunctionBuilder fb(context.getTransform());
-    // The generated tangent linkage is set to Hidden because generated tangent
-    // are never called cross-module.
-    auto linkage = SILLinkage::Hidden;
+    auto linkage =
+        jvp->isSerialized() ? SILLinkage::Public : SILLinkage::Hidden;
     auto *differential = fb.createFunction(
         linkage, diffName, diffType, diffGenericEnv, original->getLocation(),
-        original->isBare(), IsNotTransparent, original->isSerialized(),
+        original->isBare(), IsNotTransparent, jvp->isSerialized(),
         original->isDynamicallyReplaceable());
     differential->setDebugScope(
         new (module) SILDebugScope(original->getLocation(), differential));
@@ -5783,7 +5798,7 @@ bool VJPEmitter::run() {
 
 static SILFunction *createEmptyVJP(ADContext &context, SILFunction *original,
                                    SILDifferentiabilityWitness *witness,
-                                   SILLinkage linkage) {
+                                   IsSerialized_t isSerialized) {
   LLVM_DEBUG({
     auto &s = getADDebugStream();
     s << "Creating VJP:\n\t";
@@ -5817,10 +5832,10 @@ static SILFunction *createEmptyVJP(ADContext &context, SILFunction *original,
       vjpGenericSig);
 
   SILOptFunctionBuilder fb(context.getTransform());
-  auto *vjp = fb.createFunction(linkage, vjpName, vjpType, vjpGenericEnv,
-                                original->getLocation(), original->isBare(),
-                                IsNotTransparent, original->isSerialized(),
-                                original->isDynamicallyReplaceable());
+  auto *vjp = fb.createFunction(
+      witness->getLinkage(), vjpName, vjpType, vjpGenericEnv,
+      original->getLocation(), original->isBare(), IsNotTransparent,
+      isSerialized, original->isDynamicallyReplaceable());
   vjp->setDebugScope(new (module) SILDebugScope(original->getLocation(), vjp));
 
   LLVM_DEBUG(llvm::dbgs() << "VJP type: " << vjp->getLoweredFunctionType()
@@ -5830,7 +5845,7 @@ static SILFunction *createEmptyVJP(ADContext &context, SILFunction *original,
 
 static SILFunction *createEmptyJVP(ADContext &context, SILFunction *original,
                                    SILDifferentiabilityWitness *witness,
-                                   SILLinkage linkage) {
+                                   IsSerialized_t isSerialized) {
   LLVM_DEBUG({
     auto &s = getADDebugStream();
     s << "Creating JVP:\n\t";
@@ -5864,10 +5879,10 @@ static SILFunction *createEmptyJVP(ADContext &context, SILFunction *original,
       LookUpConformanceInModule(module.getSwiftModule()), jvpGenericSig);
 
   SILOptFunctionBuilder fb(context.getTransform());
-  auto *jvp = fb.createFunction(linkage, jvpName, jvpType, jvpGenericEnv,
-                                original->getLocation(), original->isBare(),
-                                IsNotTransparent, original->isSerialized(),
-                                original->isDynamicallyReplaceable());
+  auto *jvp = fb.createFunction(
+      witness->getLinkage(), jvpName, jvpType, jvpGenericEnv,
+      original->getLocation(), original->isBare(), IsNotTransparent,
+      isSerialized, original->isDynamicallyReplaceable());
   jvp->setDebugScope(new (module) SILDebugScope(original->getLocation(), jvp));
 
   LLVM_DEBUG(llvm::dbgs() << "JVP type: " << jvp->getLoweredFunctionType()
@@ -5878,7 +5893,7 @@ static SILFunction *createEmptyJVP(ADContext &context, SILFunction *original,
 /// Returns true on error.
 bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     SILFunction *original, SILDifferentiabilityWitness *witness,
-    DifferentiationInvoker invoker, bool explicitDifferentiable) {
+    DifferentiationInvoker invoker, IsSerialized_t serializeFunctions) {
   std::string traceMessage;
   llvm::raw_string_ostream OS(traceMessage);
   OS << "processing ";
@@ -5889,9 +5904,6 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
 
   assert(witness->isDefinition());
 
-  auto derivativeFunctionLinkage =
-      explicitDifferentiable ? original->getLinkage() : SILLinkage::Hidden;
-
   // If the JVP doesn't exist, need to synthesize it.
   if (!witness->getJVP()) {
     // Diagnose:
@@ -5903,7 +5915,7 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
       return true;
 
     witness->setJVP(
-        createEmptyJVP(context, original, witness, derivativeFunctionLinkage));
+        createEmptyJVP(context, original, witness, serializeFunctions));
     context.recordGeneratedFunction(witness->getJVP());
 
     // For now, only do JVP generation if the flag is enabled and if custom VJP
@@ -5974,7 +5986,7 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
       return true;
 
     witness->setVJP(
-        createEmptyVJP(context, original, witness, derivativeFunctionLinkage));
+        createEmptyVJP(context, original, witness, serializeFunctions));
     context.recordGeneratedFunction(witness->getVJP());
     VJPEmitter emitter(context, original, witness, witness->getVJP(), invoker);
     return emitter.run();
@@ -6731,7 +6743,7 @@ void Differentiation::run() {
     auto invoker = invokerPair.second;
 
     if (transformer.canonicalizeDifferentiabilityWitness(
-            original, witness, invoker, /*explicitDifferentiable*/ true))
+            original, witness, invoker, original->isSerialized()))
       errorOccurred = true;
   }
 

lib/Serialization/DeserializeSIL.cpp

@@ -3534,8 +3534,9 @@ SILDeserializer::readDifferentiabilityWitness(DeclID DId) {
     assert(!isSerialized && "declaration must not be serialized");
   }
 
-  auto linkage = fromStableSILLinkage(rawLinkage);
-  assert(linkage && "Expected value linkage for sil_differentiability_witness");
+  auto linkageOpt = fromStableSILLinkage(rawLinkage);
+  assert(linkageOpt &&
+         "Expected value linkage for sil_differentiability_witness");
   auto originalName = MF->getIdentifierText(originalNameId);
   auto jvpName = MF->getIdentifierText(jvpNameId);
   auto vjpName = MF->getIdentifierText(vjpNameId);
@@ -3572,10 +3573,14 @@ SILDeserializer::readDifferentiabilityWitness(DeclID DId) {
   auto *diffWitness =
       SILMod.lookUpDifferentiabilityWitness({originalName, config});
 
+  // Witnesses that we deserialize are always available externally; we never
+  // want to emit them ourselves.
+  auto linkage = swift::addExternalToLinkage(*linkageOpt);
+
   // If there is no existing differentiability witness, create one.
   if (!diffWitness)
     diffWitness = SILDifferentiabilityWitness::createDeclaration(
-        SILMod, *linkage, original, parameterIndices, resultIndices,
+        SILMod, linkage, original, parameterIndices, resultIndices,
         derivativeGenSig);
 
   // If the current differentiability witness is merely a declaration, and the

lib/TBDGen/TBDGen.cpp

@@ -178,6 +178,30 @@ void TBDGenVisitor::addConformances(DeclContext *DC) {
 }
 
 // SWIFT_ENABLE_TENSORFLOW
+void TBDGenVisitor::addAutoDiffLinearMapFunction(AbstractFunctionDecl *original,
+                                                 const DifferentiableAttr *attr,
+                                                 AutoDiffLinearMapKind kind) {
+  auto declRef = SILDeclRef(original);
+
+  // Linear maps are only public when the original function is serialized.
+  if (!declRef.isSerialized())
+    return;
+
+  // Differentials are only emitted when forward mode is turned on.
+  if (kind == AutoDiffLinearMapKind::Differential &&
+      !original->getASTContext()
+           .LangOpts.EnableExperimentalForwardModeDifferentiation)
+    return;
+
+  auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
+      attr->getParameterIndices(),
+      original->getInterfaceType()->castTo<AnyFunctionType>());
+  Mangle::ASTMangler mangler;
+  std::string linearMapName = mangler.mangleAutoDiffLinearMapHelper(
+      declRef.mangle(), kind, SILAutoDiffIndices(0, loweredParamIndices));
+  addSymbol(linearMapName);
+}
+
 void TBDGenVisitor::addAutoDiffDerivativeFunction(
     AbstractFunctionDecl *original, const DifferentiableAttr *attr,
     AutoDiffDerivativeFunctionKind kind) {
@@ -208,6 +232,9 @@ void TBDGenVisitor::addDifferentiabilityWitness(
 
 void TBDGenVisitor::addDifferentiableAttr(AbstractFunctionDecl *original,
                                           const DifferentiableAttr *attr) {
+  addAutoDiffLinearMapFunction(original, attr,
+                               AutoDiffLinearMapKind::Differential);
+  addAutoDiffLinearMapFunction(original, attr, AutoDiffLinearMapKind::Pullback);
   addAutoDiffDerivativeFunction(original, attr,
                                 AutoDiffDerivativeFunctionKind::JVP);
   addAutoDiffDerivativeFunction(original, attr,

lib/TBDGen/TBDGenVisitor.h

@@ -75,6 +75,12 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   void addBaseConformanceDescriptor(BaseConformance conformance);
 
   // SWIFT_ENABLE_TENSORFLOW
+  /// Adds the symbol for the linear map function of the given kind associated
+  /// with the given original function and `@differentiable` attr.
+  void addAutoDiffLinearMapFunction(AbstractFunctionDecl *original,
+                                    const DifferentiableAttr *attr,
+                                    AutoDiffLinearMapKind kind);
+
   /// Adds the symbol for the autodiff function of the given kind associated
   /// with the given original function and `@differentiable` attr.
   void addAutoDiffDerivativeFunction(AbstractFunctionDecl *original,

test/AutoDiff/Inputs/e2e_cross_module_external_module.swift

@@ -0,0 +1,14 @@
+import DifferentiationUnittest
+
+@differentiable
+public func doubleThenApplyDefaultF(_ x: Tracked<Float>) -> Tracked<Float> {
+  return x
+}
+
+@differentiable
+public func doubleThenApply(
+  _ x: Tracked<Float>,
+  _ f: @differentiable (Tracked<Float>) -> Tracked<Float> = doubleThenApplyDefaultF
+) -> Tracked<Float> {
+  return f(2 * x)
+}

test/AutoDiff/control_flow_sil.swift

@@ -69,7 +69,7 @@ func cond(_ x: Float) -> Float {
 // CHECK-SIL:   return [[VJP_RESULT]]
 
 
-// CHECK-SIL-LABEL: sil hidden [ossa] @AD__cond__pullback_src_0_wrt_0 : $@convention(thin) (Float, @owned _AD__cond_bb3__PB__src_0_wrt_0) -> Float {
+// CHECK-SIL-LABEL: sil private [ossa] @AD__cond__pullback_src_0_wrt_0 : $@convention(thin) (Float, @owned _AD__cond_bb3__PB__src_0_wrt_0) -> Float {
 // CHECK-SIL: bb0([[SEED:%.*]] : $Float, [[BB3_PB_STRUCT:%.*]] : @owned $_AD__cond_bb3__PB__src_0_wrt_0):
 // CHECK-SIL:   [[BB3_PRED:%.*]] = destructure_struct [[BB3_PB_STRUCT]] : $_AD__cond_bb3__PB__src_0_wrt_0
 // CHECK-SIL:   switch_enum [[BB3_PRED]] : $_AD__cond_bb3__Pred__src_0_wrt_0, case #_AD__cond_bb3__Pred__src_0_wrt_0.bb2!enumelt.1: bb3, case #_AD__cond_bb3__Pred__src_0_wrt_0.bb1!enumelt.1: bb1
@@ -217,7 +217,7 @@ func cond_tuple_var(_ x: Float) -> Float {
   return y.1
 }
 
-// CHECK-SIL-LABEL: sil hidden [ossa] @AD__cond_tuple_var__pullback_src_0_wrt_0 : $@convention(thin) (Float, @owned _AD__cond_tuple_var_bb3__PB__src_0_wrt_0) -> Float {
+// CHECK-SIL-LABEL: sil private [ossa] @AD__cond_tuple_var__pullback_src_0_wrt_0 : $@convention(thin) (Float, @owned _AD__cond_tuple_var_bb3__PB__src_0_wrt_0) -> Float {
 // CHECK-SIL: bb0([[SEED:%.*]] : $Float, [[BB3_PB_STRUCT:%.*]] : $_AD__cond_tuple_var_bb3__PB__src_0_wrt_0):
 // CHECK-SIL:   [[BB3_PRED:%.*]] = destructure_struct [[BB3_PB_STRUCT]] : $_AD__cond_tuple_var_bb3__PB__src_0_wrt_0
 // CHECK-SIL:   copy_addr {{%.*}} to {{%.*}} : $*(Float, Float)

test/AutoDiff/differentiation_transform_diagnostics.swift

@@ -466,3 +466,35 @@ _ = gradient(at: Float(1), Float(2), in: (+) as @differentiable (Float, @nondiff
 
 // expected-error @+1 {{conversion to '@differentiable(linear)' function type is not yet supported}}
 let _: @differentiable(linear) (Float) -> Float = { x in x }
+
+//===----------------------------------------------------------------------===//
+// Differentiating from fragile functions
+//===----------------------------------------------------------------------===//
+
+public func implicitlyDifferentiableFromFragile(_ x: Float) -> Float { x }
+
+public func hasImplicitlyDifferentiatedTopLevelDefaultArgument(
+// expected-error @+2 {{function is not differentiable}}
+// expected-note @+1 {{differentiated functions in default arguments must be marked '@differentiable' or have a public '@derivative'}}
+  _ f: @differentiable (Float) -> Float = implicitlyDifferentiableFromFragile
+) {}
+
+// TODO(TF-1030): This will eventually not be an error.
+// expected-error @+2 {{function is not differentiable}}
+// expected-note @+1 {{differentiated functions in default arguments must be marked '@differentiable' or have a public '@derivative'; this is not possible with a closure, make a top-level function instead}}
+public func hasImplicitlyDifferentiatedClosureDefaultArgument(_ f: @differentiable (Float) -> Float = { $0 }) {}
+
+@inlinable
+public func fragileFuncWithGradient() {
+  // expected-error @+2 {{function is not differentiable}}
+  // expected-note @+1 {{differentiated functions in '@inlinable' functions must be marked '@differentiable' or have a public '@derivative'}}
+  let _ = gradient(at: 0, in: implicitlyDifferentiableFromFragile)
+}
+
+@inlinable
+@differentiable
+public func fragileDifferentiable(_ x: Float) -> Float {
+  // expected-error @+2 {{expression is not differentiable}}
+  // expected-note @+1 {{differentiated functions in '@inlinable' functions must be marked '@differentiable' or have a public '@derivative'}}
+  implicitlyDifferentiableFromFragile(x)
+}