[AutoDiff] SILGen differentiability witnesses.

Generate SIL differentiability witnesses from AST `@differentiable` and
`@differentiating` attributes.

Author: dan-zheng

include/swift/AST/Attr.h

@@ -1682,6 +1682,8 @@ class DifferentiatingAttr final
   unsigned NumParsedParameters = 0;
   /// The differentiation parameters' indices, resolved by the type checker.
   IndexSubset *ParameterIndices = nullptr;
+  /// The derivative function kind (JVP or VJP), resolved by the type checker.
+  Optional<AutoDiffDerivativeFunctionKind> Kind = None;
 
   explicit DifferentiatingAttr(ASTContext &context, bool implicit,
                                SourceLoc atLoc, SourceRange baseRange,
@@ -1711,6 +1713,12 @@ class DifferentiatingAttr final
   FuncDecl *getOriginalFunction() const { return OriginalFunction; }
   void setOriginalFunction(FuncDecl *decl) { OriginalFunction = decl; }
 
+  AutoDiffDerivativeFunctionKind getDerivativeKind() const {
+    assert(Kind && "Derivative function kind has not yet been resolved");
+    return *Kind;
+  }
+  void setDerivativeKind(AutoDiffDerivativeFunctionKind kind) { Kind = kind; }
+
   /// The parsed differentiation parameters, i.e. the list of parameters
   /// specified in 'wrt:'.
   ArrayRef<ParsedAutoDiffParameter> getParsedParameters() const {

include/swift/AST/AutoDiff.h

@@ -413,6 +413,27 @@ template<> struct DenseMapInfo<AutoDiffConfig> {
   }
 };
 
+template<> struct DenseMapInfo<AutoDiffDerivativeFunctionKind> {
+  static AutoDiffDerivativeFunctionKind getEmptyKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getEmptyKey());
+  }
+
+  static AutoDiffDerivativeFunctionKind getTombstoneKey() {
+    return static_cast<AutoDiffDerivativeFunctionKind::innerty>(
+        DenseMapInfo<unsigned>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const AutoDiffDerivativeFunctionKind &Val) {
+    return DenseMapInfo<unsigned>::getHashValue(Val);
+  }
+
+  static bool isEqual(const AutoDiffDerivativeFunctionKind &LHS,
+                      const AutoDiffDerivativeFunctionKind &RHS) {
+    return LHS == RHS;
+  }
+};
+
 template<> struct DenseMapInfo<SILAutoDiffIndices> {
   static SILAutoDiffIndices getEmptyKey() {
     return { DenseMapInfo<unsigned>::getEmptyKey(), nullptr };

include/swift/SIL/SILModule.h

@@ -609,6 +609,13 @@ class SILModule {
   /// hierarchy of \p Class.
   SILFunction *lookUpFunctionInVTable(ClassDecl *Class, SILDeclRef Member);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Look up the differentiability witness corresponding to the given key.
+  SILDifferentiabilityWitness *
+  lookUpDifferentiabilityWitness(SILDifferentiabilityWitnessKey key,
+                                 bool deserializeLazily=true);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   // Given a protocol, attempt to create a default witness table declaration
   // for it.
   SILDefaultWitnessTable *

lib/SIL/SILModule.cpp

@@ -251,6 +251,18 @@ SILModule::lookUpDefaultWitnessTable(const ProtocolDecl *Protocol,
   return found->second;
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+SILDifferentiabilityWitness *
+SILModule::lookUpDifferentiabilityWitness(SILDifferentiabilityWitnessKey key,
+                                          bool deserializeLazily) {
+  auto found = DifferentiabilityWitnessMap.find(key);
+  if (found != DifferentiabilityWitnessMap.end())
+    return found->second;
+  if (deserializeLazily)
+    return getSILLoader()->lookupDifferentiabilityWitness(key);
+  return nullptr;
+}
+
 SILDefaultWitnessTable *
 SILModule::createDefaultWitnessTableDeclaration(const ProtocolDecl *Protocol,
                                                 SILLinkage Linkage) {

lib/SILGen/SILGen.cpp

@@ -752,87 +752,140 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
              F->print(llvm::dbgs()));
 
   // SWIFT_ENABLE_TENSORFLOW
-  // Create self-reordering thunks for JVPs/VJPs of `@differentiable` methods.
-  if (constant.hasDecl() && constant.getAbstractFunctionDecl()) {
+  // Visit `@differentiable` and `@differentiating` attributes and generate SIL
+  // differentiability witnesses.
+  // Do not visit default argument generator functions.
+  if (constant.hasDecl() && constant.getAbstractFunctionDecl() &&
+      constant.kind != SILDeclRef::Kind::DefaultArgGenerator) {
     auto *AFD = constant.getAbstractFunctionDecl();
-    auto origFnType = AFD->getInterfaceType()->castTo<AnyFunctionType>();
-    auto origSilFnType = F->getLoweredFunctionType();
-    // Jointly iterate over AST `@differentiable` attributes and SIL
-    // `[differentiable]` attributes.
-    auto diffAttrs = AFD->getAttrs().getAttributes<DifferentiableAttr>();
-    auto silDiffAttrs = F->getDifferentiableAttrs();
-    for (auto pair : llvm::zip(diffAttrs, silDiffAttrs)) {
-      auto *diffAttr = const_cast<DifferentiableAttr *>(std::get<0>(pair));
-      auto *silDiffAttr = std::get<1>(pair);
-      // Compute lowered parameter indices.
-      auto *paramIndices = diffAttr->getParameterIndices();
-      auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
-          paramIndices, origFnType);
-      SILAutoDiffIndices indices(/*source*/ 0, loweredParamIndices);
-      assert(silDiffAttr->getIndices() == indices &&
-             "Expected matching @differentiable and [differentiable] indices");
-
-      auto lookUpConformance = LookUpConformanceInModule(M.getSwiftModule());
-      auto expectedJVPType = origSilFnType->getAutoDiffDerivativeFunctionType(
-          indices.parameters, indices.source,
-          AutoDiffDerivativeFunctionKind::JVP, Types, lookUpConformance);
-      auto expectedVJPType = origSilFnType->getAutoDiffDerivativeFunctionType(
-          indices.parameters, indices.source,
-          AutoDiffDerivativeFunctionKind::VJP, Types, lookUpConformance);
-
-      // Self reordering is necessary if wrt at least two parameters, including
-      // self.
-      auto shouldReorderSelf = [&]() {
-        if (!F->hasSelfParam())
-          return false;
-        auto selfParamIndex = origSilFnType->getNumParameters() - 1;
-        if (!indices.isWrtParameter(selfParamIndex))
-          return false;
-        return indices.parameters->getNumIndices() > 1;
-      };
-      bool reorderSelf = shouldReorderSelf();
-
-      // Thunk JVP method, if it is defined.
-      if (auto *jvpDecl = diffAttr->getJVPFunction()) {
-        SILFunction *jvpThunk;
-        auto *jvpFn = getFunction(SILDeclRef(jvpDecl), NotForDefinition);
-        if (reorderSelf || jvpFn->getLoweredFunctionType() != expectedJVPType) {
-          jvpThunk = getOrCreateAutoDiffDerivativeFunctionThunk(
-              F, indices, jvpFn, AutoDiffDerivativeFunctionKind::JVP,
-              reorderSelf);
-        } else {
-          auto *id = AutoDiffDerivativeFunctionIdentifier::get(
-              AutoDiffDerivativeFunctionKind::JVP,
-              diffAttr->getParameterIndices(), AFD->getASTContext());
-          jvpThunk = getOrCreateAutoDiffThunk(
-              constant.asAutoDiffDerivativeFunction(id), jvpFn,
-              expectedJVPType);
-        }
-        silDiffAttr->setJVPName(jvpThunk->getName());
-      }
-      // Thunk VJP method, if it is defined.
-      if (auto *vjpDecl = diffAttr->getVJPFunction()) {
-        SILFunction *vjpThunk;
-        auto *vjpFn = getFunction(SILDeclRef(vjpDecl), NotForDefinition);
-        if (reorderSelf || vjpFn->getLoweredFunctionType() != expectedVJPType) {
-          vjpThunk = getOrCreateAutoDiffDerivativeFunctionThunk(
-              F, indices, vjpFn, AutoDiffDerivativeFunctionKind::VJP,
-              reorderSelf);
-        } else {
-          auto *id = AutoDiffDerivativeFunctionIdentifier::get(
-              AutoDiffDerivativeFunctionKind::VJP,
-              diffAttr->getParameterIndices(), AFD->getASTContext());
-          vjpThunk = getOrCreateAutoDiffThunk(
-              constant.asAutoDiffDerivativeFunction(id), vjpFn,
-              expectedVJPType);
-        }
-        silDiffAttr->setVJPName(vjpThunk->getName());
+    // 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);
+      emitDifferentiabilityWitness(AFD, F, diffAttr->getParameterIndices(), jvp,
+                                   vjp);
+    }
+    // Visit all `@differentiating` attributes.
+    for (auto *diffAttr :
+             AFD->getAttrs().getAttributes<DifferentiatingAttr>()) {
+      auto *origAFD = diffAttr->getOriginalFunction();
+      auto *origFn = getFunction(SILDeclRef(origAFD), NotForDefinition);
+      SILFunction *jvp = nullptr;
+      SILFunction *vjp = nullptr;
+      switch (diffAttr->getDerivativeKind()) {
+      case AutoDiffDerivativeFunctionKind::JVP:
+        jvp = F;
+        break;
+      case AutoDiffDerivativeFunctionKind::VJP:
+        vjp = F;
+        break;
       }
+      emitDifferentiabilityWitness(origAFD, origFn,
+                                   diffAttr->getParameterIndices(), jvp, vjp);
     }
   }
   F->verify();
 }
 
+void SILGenModule::emitDifferentiabilityWitness(
+    AbstractFunctionDecl *originalAFD, SILFunction *originalFunction,
+    IndexSubset *parameterIndices, SILFunction *jvp, SILFunction *vjp) {
+  auto *origFnType = originalAFD->getInterfaceType()->castTo<AnyFunctionType>();
+  auto origSilFnType = originalFunction->getLoweredFunctionType();
+  auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
+      parameterIndices, origFnType);
+  // NOTE(TF-893): Extending capacity is necessary when `origSilFnType` has
+  // parameters corresponding to captured variables. These parameters do not
+  // appear in the type of `origFnType`.
+  // TODO: If posssible, change `autodiff::getLoweredParameterIndices` to
+  // take `CaptureInfo` into account.
+  if (origSilFnType->getNumParameters() > loweredParamIndices->getCapacity())
+    loweredParamIndices = loweredParamIndices->extendingCapacity(
+        getASTContext(), origSilFnType->getNumParameters());
+  SILAutoDiffIndices indices(/*source*/ 0, loweredParamIndices);
+
+  // Self reordering thunk is necessary if wrt at least two parameters,
+  // including self.
+  auto shouldReorderSelf = [&]() {
+    if (!originalFunction->hasSelfParam())
+      return false;
+    auto selfParamIndex = origSilFnType->getNumParameters() - 1;
+    if (!indices.isWrtParameter(selfParamIndex))
+      return false;
+    return indices.parameters->getNumIndices() > 1;
+  };
+  bool reorderSelf = shouldReorderSelf();
+
+  // Get or create differentiability witness.
+  CanGenericSignature derivativeGenSig;
+  if (jvp && vjp)
+    assert(jvp->getLoweredFunctionType()->getGenericSignature() ==
+               vjp->getLoweredFunctionType()->getGenericSignature() &&
+           "JVP and VJP generic signatures must match");
+  if (jvp)
+    derivativeGenSig = jvp->getLoweredFunctionType()->getGenericSignature();
+  if (vjp)
+    derivativeGenSig = vjp->getLoweredFunctionType()->getGenericSignature();
+  auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
+  AutoDiffConfig config{loweredParamIndices, resultIndices,
+                        derivativeGenSig};
+  auto key = std::make_pair(originalFunction->getName(), config);
+  SILDifferentiabilityWitness *diffWitness = nullptr;
+  if (auto *foundWitness = M.lookUpDifferentiabilityWitness(
+          key, /*deserializeLazily*/ false)) {
+    diffWitness = foundWitness;
+  } else {
+    // Create new SIL differentiability witness.
+    diffWitness = SILDifferentiabilityWitness::create(
+        M, originalFunction->getLinkage(), originalFunction,
+        loweredParamIndices, resultIndices, derivativeGenSig, /*jvp*/ nullptr,
+        /*vjp*/ nullptr, /*isSerialized*/ true);
+  }
+
+  // Set derivative function in differentiability witness.
+  auto setDerivativeInDifferentiabilityWitness =
+      [&](AutoDiffDerivativeFunctionKind kind, SILFunction *derivative) {
+    auto expectedDerivativeType =
+        origSilFnType->getAutoDiffDerivativeFunctionType(
+            indices.parameters, indices.source, kind, Types,
+            LookUpConformanceInModule(M.getSwiftModule()));
+    // Thunk derivative function.
+    SILFunction *derivativeThunk;
+    if (reorderSelf ||
+        derivative->getLoweredFunctionType() != expectedDerivativeType) {
+      derivativeThunk = getOrCreateAutoDiffDerivativeFunctionThunk(
+          originalFunction, indices, derivative, kind, reorderSelf);
+    } else {
+      auto *id = AutoDiffDerivativeFunctionIdentifier::get(
+          kind, parameterIndices, getASTContext());
+      derivativeThunk = getOrCreateAutoDiffThunk(
+          SILDeclRef(originalAFD).asAutoDiffDerivativeFunction(id), derivative,
+          expectedDerivativeType);
+    }
+    // Check for existing same derivative.
+    // TODO(TF-898): Remove condition below and simplify assertion to
+    // `!diffWitness->getDerivative(kind)` after `@differentiating` attribute
+    // type-checking no longer generates implicit `@differentiable` attributes.
+    auto *existingDerivative = diffWitness->getDerivative(kind);
+    if (existingDerivative && existingDerivative == derivativeThunk)
+      return;
+    assert(!existingDerivative &&
+           "SIL differentiability witness already has a different existing "
+           "derivative");
+    diffWitness->setDerivative(kind, derivativeThunk);
+  };
+  if (jvp)
+    setDerivativeInDifferentiabilityWitness(AutoDiffDerivativeFunctionKind::JVP,
+                                            jvp);
+  if (vjp)
+    setDerivativeInDifferentiabilityWitness(AutoDiffDerivativeFunctionKind::VJP,
+                                            vjp);
+}
+
 void SILGenModule::
 emitMarkFunctionEscapeForTopLevelCodeGlobals(SILLocation loc,
                                              const CaptureInfo &captureInfo) {

lib/SILGen/SILGen.h

@@ -318,6 +318,16 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Emit the self-conformance witness table for a protocol.
   void emitSelfConformanceWitnessTable(ProtocolDecl *protocol);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Emit the differentiability witness for the given original function
+  /// declaration and SIL function, parameter indices, and JVP and VJP
+  /// functions (null if undefined).
+  void emitDifferentiabilityWitness(AbstractFunctionDecl *originalAFD,
+                                    SILFunction *originalFunction,
+                                    IndexSubset *parameterIndices,
+                                    SILFunction *jvp, SILFunction *vjp);
+  // SWIFT_ENABLE_TENSORFLOW END
+
   /// Emit the lazy initializer function for a global pattern binding
   /// declaration.
   SILFunction *emitLazyGlobalInitializer(StringRef funcName,

lib/Sema/TypeCheckAttr.cpp

@@ -3610,6 +3610,7 @@ void AttributeChecker::visitDifferentiatingAttr(DifferentiatingAttr *attr) {
     attr->setInvalid();
     return;
   }
+  attr->setDerivativeKind(kind);
   // `value: R` result tuple element must conform to `Differentiable`.
   auto diffableProto = ctx.getProtocol(KnownProtocolKind::Differentiable);
   auto valueResultType = valueResultElt.getType();

lib/Serialization/DeserializeSIL.cpp

@@ -2972,6 +2972,9 @@ void SILDeserializer::readWitnessTableEntries(
   // Another record means the end of this WitnessTable.
   while (kind != SIL_WITNESS_TABLE &&
          kind != SIL_DEFAULT_WITNESS_TABLE &&
+         // SWIFT_ENABLE_TENSORFLOW
+         kind != SIL_DIFFERENTIABILITY_WITNESS &&
+         // SWIFT_ENABLE_TENSORFLOW END
          kind != SIL_FUNCTION) {
     if (kind == SIL_DEFAULT_WITNESS_TABLE_NO_ENTRY) {
       witnessEntries.push_back(SILDefaultWitnessTable::Entry());

test/AutoDiff/sil_differentiability_witness_silgen.swift

@@ -0,0 +1,54 @@
+// RUN: %target-swift-frontend -emit-silgen %s | %target-sil-opt | %FileCheck %s
+
+// Test SIL differentiability witness SIL generation.
+
+// Test public non-generic function.
+// SIL differentiability witness:
+// - Has public linkage (implicit).
+// - Has no `where` clause.
+
+public func foo(_ x: Float) -> Float { x }
+
+@differentiating(foo)
+public func foo_jvp(_ x: Float) -> (value: Float, differential: (Float) -> Float) {
+  (x, { $0 })
+}
+
+@differentiating(foo)
+public func foo_vjp(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { $0 })
+}
+
+// CHECK-LABEL: // differentiability witness for foo(_:)
+// CHECK-NEXT: sil_differentiability_witness [serialized] [parameters 0] [results 0] @$s36sil_differentiability_witness_silgen3fooyS2fF : $@convention(thin) (Float) -> Float {
+// CHECK-NEXT:   jvp: @AD__$s36sil_differentiability_witness_silgen3fooyS2fF__jvp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-NEXT:   vjp: @AD__$s36sil_differentiability_witness_silgen3fooyS2fF__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-NEXT: }
+
+// Test internal generic function.
+// SIL differentiability witness:
+// - Has hidden linkage.
+// - Has `where` clause.
+
+@differentiable(where T: Differentiable)
+func generic<T>(_ x: T, _ y: Float) -> T { x }
+
+@differentiating(generic)
+func generic_jvp<T: Differentiable>(_ x: T, _ y: Float) -> (
+  value: T, differential: (T.TangentVector, Float) -> T.TangentVector
+) {
+  (x, { dx, dy in dx })
+}
+
+@differentiating(generic)
+func generic_vjp<T: Differentiable>(_ x: T, _ y: Float) -> (
+  value: T, pullback: (T.TangentVector) -> (T.TangentVector, Float)
+) {
+  (x, { ($0, .zero) })
+}
+
+// CHECK-LABEL: // differentiability witness for generic<A>(_:_:)
+// CHECK-NEXT: sil_differentiability_witness hidden [serialized] [parameters 0 1] [results 0] [where τ_0_0 : _Differentiable] @$s36sil_differentiability_witness_silgen7genericyxx_SftlF : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0, Float) -> @out τ_0_0 {
+// CHECK-NEXT:   jvp: @AD__$s36sil_differentiability_witness_silgen7genericyxx_SftlF__jvp_src_0_wrt_0_1 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector, Float) -> @out τ_0_0.TangentVector)
+// CHECK-NEXT:   vjp: @AD__$s36sil_differentiability_witness_silgen7genericyxx_SftlF__vjp_src_0_wrt_0_1 : $@convention(thin) <τ_0_0 where τ_0_0 : _Differentiable> (@in_guaranteed τ_0_0, Float) -> (@out τ_0_0, @owned @callee_guaranteed (@in_guaranteed τ_0_0.TangentVector) -> (@out τ_0_0.TangentVector, Float))
+// CHECK-NEXT: }