[AutoDiff] Fix two issues related with emission of differentiability witnesses (#80983)

1. When differentiable nested function (closure) is specialized by capture promotion pass ensure we generate a differentiability witness for the specialized function as well. Ensure the original witness is removed if the original function becomes dead. 
2. Differentiability witnesses for a function could originate either from its `@differentiable` attribute or from explicit `@derivative(of:)` attribute on the derivative. In the latter case the derivative itself might not be emitted, while original function is (e.g. original function is `@inlineable`, but derivative is `@usableFromInline`). Previously both cases were handled only when function body was emitted. As a result we missed witness in the aforementioned case. Ensure the
differentiability witness originating from `@derivative(of:)` is emitted even if we're not going to emit body of the derivative.

Fixes #59135

Author: asl

include/swift/SIL/SILModule.h

@@ -591,6 +591,12 @@ class SILModule {
   /// Erase a global SIL variable from the module.
   void eraseGlobalVariable(SILGlobalVariable *G);
 
+  /// Erase a differentiability witness from the module.
+  void eraseDifferentiabilityWitness(SILDifferentiabilityWitness *dw);
+
+  /// Erase all differentiability witnesses for function f.
+  void eraseAllDifferentiabilityWitnesses(SILFunction *f);
+
   /// Create and return an empty SIL module suitable for generating or parsing
   /// SIL into.
   ///

lib/SIL/IR/SILModule.cpp

@@ -504,6 +504,33 @@ void SILModule::eraseGlobalVariable(SILGlobalVariable *gv) {
   getSILGlobalList().erase(gv);
 }
 
+void SILModule::eraseDifferentiabilityWitness(SILDifferentiabilityWitness *dw) {
+  getSILLoader()->invalidateDifferentiabilityWitness(dw);
+
+  Mangle::ASTMangler mangler(getASTContext());
+  auto originalFunction = dw->getOriginalFunction()->getName();
+  auto mangledKey = mangler.mangleSILDifferentiabilityWitness(
+    originalFunction, dw->getKind(), dw->getConfig());
+  DifferentiabilityWitnessMap.erase(mangledKey);
+  llvm::erase(DifferentiabilityWitnessesByFunction[originalFunction], dw);
+
+  getDifferentiabilityWitnessList().erase(dw);
+}
+
+void SILModule::eraseAllDifferentiabilityWitnesses(SILFunction *f) {
+  Mangle::ASTMangler mangler(getASTContext());
+
+  for (auto *dw : DifferentiabilityWitnessesByFunction.at(f->getName())) {
+    getSILLoader()->invalidateDifferentiabilityWitness(dw);
+    auto mangledKey = mangler.mangleSILDifferentiabilityWitness(
+      f->getName(), dw->getKind(), dw->getConfig());
+    DifferentiabilityWitnessMap.erase(mangledKey);
+    getDifferentiabilityWitnessList().erase(dw);
+  }
+
+  DifferentiabilityWitnessesByFunction.erase(f->getName());
+}
+
 SILVTable *SILModule::lookUpVTable(const ClassDecl *C,
                                    bool deserializeLazily) {
   if (!C)

lib/SILGen/SILGen.cpp

@@ -1409,7 +1409,7 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
 
 void SILGenModule::emitDifferentiabilityWitnessesForFunction(
     SILDeclRef constant, SILFunction *F) {
-  // Visit `@derivative` attributes and generate SIL differentiability
+  // Visit `@differentiable` attributes and generate SIL differentiability
   // witnesses.
   // Skip if the SILDeclRef is a:
   // - Default argument generator function.
@@ -1439,33 +1439,6 @@ void SILGenModule::emitDifferentiabilityWitnessesForFunction(
                                    config, /*jvp*/ nullptr,
                                    /*vjp*/ nullptr, diffAttr);
     }
-    for (auto *derivAttr : Attrs.getAttributes<DerivativeAttr>()) {
-      SILFunction *jvp = nullptr;
-      SILFunction *vjp = nullptr;
-      switch (derivAttr->getDerivativeKind()) {
-      case AutoDiffDerivativeFunctionKind::JVP:
-        jvp = F;
-        break;
-      case AutoDiffDerivativeFunctionKind::VJP:
-        vjp = F;
-        break;
-      }
-      auto *origAFD = derivAttr->getOriginalFunction(getASTContext());
-      auto origDeclRef =
-          SILDeclRef(origAFD).asForeign(requiresForeignEntryPoint(origAFD));
-      auto *origFn = getFunction(origDeclRef, NotForDefinition);
-      auto witnessGenSig =
-          autodiff::getDifferentiabilityWitnessGenericSignature(
-              origAFD->getGenericSignature(), AFD->getGenericSignature());
-      auto *resultIndices =
-        autodiff::getFunctionSemanticResultIndices(origAFD,
-                                                   derivAttr->getParameterIndices());
-      AutoDiffConfig config(derivAttr->getParameterIndices(), resultIndices,
-                            witnessGenSig);
-      emitDifferentiabilityWitness(origAFD, origFn,
-                                   DifferentiabilityKind::Reverse, config, jvp,
-                                   vjp, derivAttr);
-    }
   };
   if (auto *accessor = dyn_cast<AccessorDecl>(AFD))
     if (accessor->isGetter())
@@ -1582,6 +1555,43 @@ void SILGenModule::emitAbstractFuncDecl(AbstractFunctionDecl *AFD) {
         SILDeclRef::BackDeploymentKind::Thunk);
     emitBackDeploymentThunk(thunk);
   }
+
+  // Emit differentiability witness for the function referenced in
+  // @derivative(of:) attribute registering current function as VJP / JVP.
+  // Differentiability witnesses for a function could originate either from its
+  // @differentiable attribute or from explicit @derivative(of:) attribute on
+  // the derivative. In the latter case the derivative itself might not be
+  // emitted, while original function is (e.g. original function is @inlineable,
+  // but derivative is @usableFromInline). Ensure the differentiability witness
+  // originating from @derivative(of:) is emitted even if we're not going to
+  // emit body of the derivative.
+  for (auto *derivAttr : AFD->getAttrs().getAttributes<DerivativeAttr>()) {
+      auto *f = getFunction(SILDeclRef(AFD), NotForDefinition);
+      SILFunction *jvp = nullptr, *vjp = nullptr;
+      switch (derivAttr->getDerivativeKind()) {
+      case AutoDiffDerivativeFunctionKind::JVP:
+        jvp = f;
+        break;
+      case AutoDiffDerivativeFunctionKind::VJP:
+        vjp = f;
+        break;
+      }
+      auto *origAFD = derivAttr->getOriginalFunction(getASTContext());
+      auto origDeclRef =
+          SILDeclRef(origAFD).asForeign(requiresForeignEntryPoint(origAFD));
+      auto *origFn = getFunction(origDeclRef, NotForDefinition);
+      auto witnessGenSig =
+          autodiff::getDifferentiabilityWitnessGenericSignature(
+              origAFD->getGenericSignature(), AFD->getGenericSignature());
+      auto *resultIndices =
+        autodiff::getFunctionSemanticResultIndices(origAFD,
+                                                   derivAttr->getParameterIndices());
+      AutoDiffConfig config(derivAttr->getParameterIndices(), resultIndices,
+                            witnessGenSig);
+      emitDifferentiabilityWitness(origAFD, origFn,
+                                   DifferentiabilityKind::Reverse, config, jvp,
+                                   vjp, derivAttr);
+    }
 }
 
 void SILGenModule::emitFunction(FuncDecl *fd) {

lib/SILOptimizer/Mandatory/CapturePromotion.cpp

@@ -1473,9 +1473,25 @@ processPartialApplyInst(SILOptFunctionBuilder &funcBuilder,
       funcBuilder, pai, fri, promotableIndices, f->getResilienceExpansion());
   worklist.push_back(clonedFn);
 
+  SILFunction *origFn = fri->getReferencedFunction();
+  for (const auto *w : mod.lookUpDifferentiabilityWitnessesForFunction(
+         origFn->getName())) {
+    // @derivative(of:) attribute could only be applied at global scope, therefore
+    // local functions might not have custom derivatives registered
+    assert(!w->getJVP() && !w->getVJP() && "does not expect custom derivatives here");
+    auto linkage = stripExternalFromLinkage(clonedFn->getLinkage());
+    SILDifferentiabilityWitness::createDefinition(
+      mod, linkage, clonedFn,
+      w->getKind(), w->getParameterIndices(), w->getResultIndices(),
+      w->getDerivativeGenericSignature(),
+      /*jvp*/ nullptr, /*vjp*/ nullptr,
+      /*isSerialized*/ hasPublicVisibility(clonedFn->getLinkage()),
+      w->getAttribute());
+  }
+
   // Mark the original partial apply function as deletable if it doesn't have
   // uses later.
-  fri->getReferencedFunction()->addSemanticsAttr(semantics::DELETE_IF_UNUSED);
+  origFn->addSemanticsAttr(semantics::DELETE_IF_UNUSED);
 
   // Initialize a SILBuilder and create a function_ref referencing the cloned
   // closure.

lib/SILOptimizer/Mandatory/DiagnosticDeadFunctionElimination.cpp

@@ -36,6 +36,7 @@ namespace {
 struct DiagnosticDeadFunctionEliminator : SILFunctionTransform {
   void run() override {
     auto *fn = getFunction();
+    auto &mod = fn->getModule();
 
     // If an earlier pass asked us to eliminate the function body if it's
     // unused, and the function is in fact unused, do that now.
@@ -67,6 +68,10 @@ struct DiagnosticDeadFunctionEliminator : SILFunctionTransform {
       b.createUnreachable(loc);
     }
 
+    // Drop differentiability witnesses, if any
+    if (!mod.lookUpDifferentiabilityWitnessesForFunction(fn->getName()).empty())
+      mod.eraseAllDifferentiabilityWitnesses(fn);
+
     // If the function has shared linkage, reduce this version to private
     // linkage, because we don't want the deleted-body form to win in any
     // ODR shootouts.

test/AutoDiff/compiler_crashers_fixed/issue-59135-nested-function-diff-wittness-capture-promotion.swift

@@ -0,0 +1,33 @@
+// RUN: %target-swift-emit-sil -Xllvm -debug-only=differentiation -o /dev/null 2>&1 %s | %FileCheck %s
+
+// The differentiability witness for y in s(h:) will be generated by silgen. However, later the capture
+// promotion pass would specialize it since it only captures an integer and therefore does not need to
+// box the capture. Ensure we create differentiability witness for specialized function. In addition to
+// this, since the original function is not used anymore, the body of it is removed (with only unreachable
+// terminator inside). Remove original differentiability witness as it would lead to non-differentiable
+// diagnostics further on.
+
+// CHECK-LABEL: differentiability witness for specialized y #1 (_:) in s(h:)
+// CHECK: sil_differentiability_witness private [reverse] [parameters 0] [results 0] @$s{{.*}} : $@convention(thin) (@guaranteed W, Int) -> @owned W {
+// CHECK-NOT: sil_differentiability_witness private [reverse] [parameters 0] [results 0] @$s{{.*}} : $@convention(thin) (@guaranteed W, @guaranteed { var Int }) -> @owned W {
+
+import _Differentiation
+struct B: Differentiable{}
+struct X { var j = [Float]()}
+struct W: Differentiable {
+    @noDerivative var z: X
+    var h: B
+}
+func o<T, R>(_ x: T, _ f: @differentiable(reverse) (T) -> R) -> R {f(x)}
+func m<T, R>(_ f: @escaping @differentiable(reverse) (T) -> R) -> @differentiable(reverse) (T) -> R {{ x in o(x, f) }}
+@differentiable(reverse)
+ func s(h: B) -> B {
+    var (_, e) = (0,0)
+    @differentiable(reverse)
+    func y(_ i: W) -> W {
+        let _ = e;
+        return i
+    }
+    let w = m(y)
+    return B()
+}

test/AutoDiff/compiler_crashers_fixed/issue-59135-usableFromInline-VJP.swift

@@ -0,0 +1,38 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-emit-sil -emit-module -module-name M -emit-module-path %t/M.swiftmodule 2>&1 %s | %FileCheck %s
+
+// The original function Tensor.subscriptIndexPath() is not marked as @differentiable. As a result, no explicit differentiable witness is generated for it.
+// However, the witness is generated as a side effect of providing a derivative via  @derivative(of: subscriptIndexPath) on _vjpSubscriptIndexPath.
+// Since _vjpSubscriptIndexPath is not emitted when -emit-module is used, we need to ensure we still generate a witness.
+
+import _Differentiation
+
+// CHECK-LABEL: differentiability witness for Tensor.subscriptIndexPath()
+// CHECK: sil_differentiability_witness [serialized] [reverse] [parameters 0] [results 0] @$s1M6TensorV18subscriptIndexPathACyF : $@convention(method) (Tensor) -> Tensor {
+// CHECK:  vjp: @$s1M6TensorV18subscriptIndexPathACyFTJrSpSr : $@convention(method) (Tensor) -> (Tensor, @owned @callee_guaranteed (Tensor) -> Tensor)
+
+// CHECK-LABEL: reverse-mode derivative of Tensor.subscriptIndexPath()
+// CHECK: @$s1M6TensorV18subscriptIndexPathACyFTJrSpSr : $@convention(method) (Tensor) -> (Tensor, @owned @callee_guaranteed (Tensor) -> Tensor) {
+// CHECK: function_ref Tensor._vjpSubscriptIndexPath()
+// CHECK: function_ref @$s1M6TensorV22_vjpSubscriptIndexPathAC5value_A2Cc8pullbacktyF : $@convention(method) (Tensor) -> (Tensor, @owned @callee_guaranteed (Tensor) -> Tensor)
+
+public struct Tensor: Differentiable & AdditiveArithmetic {
+  @inlinable
+  func subscriptIndexPath() -> Tensor {
+    fatalError()
+  }
+
+  @inlinable
+  @differentiable(reverse, wrt: self)
+  func subscriptRanges() -> Tensor {
+    subscriptIndexPath()
+  }
+  
+  @usableFromInline
+  @derivative(of: subscriptIndexPath)
+  func _vjpSubscriptIndexPath() -> (
+    value: Tensor, pullback: (Tensor) -> Tensor
+  ) {
+    fatalError()
+  }
+}