[AutoDiff upstream] Add derivative function witness/vtable entry SILGen.

include/swift/AST/AutoDiff.h

@@ -104,9 +104,8 @@ class AutoDiffDerivativeFunctionIdentifier : public llvm::FoldingSetNode {
   void Profile(llvm::FoldingSetNodeID &ID) {
     ID.AddInteger(kind);
     ID.AddPointer(parameterIndices);
-    CanGenericSignature derivativeCanGenSig;
-    if (derivativeGenericSignature)
-      derivativeCanGenSig = derivativeGenericSignature->getCanonicalSignature();
+    auto derivativeCanGenSig =
+        derivativeGenericSignature.getCanonicalSignature();
     ID.AddPointer(derivativeCanGenSig.getPointer());
   }
 };

include/swift/SIL/SILDeclRef.h

@@ -324,6 +324,13 @@ struct SILDeclRef {
     return declRef;
   }
 
+  /// Returns this `SILDeclRef` replacing `loc` with `decl`.
+  SILDeclRef withDecl(ValueDecl *decl) const {
+    SILDeclRef result = *this;
+    result.loc = decl;
+    return result;
+  }
+
   /// True if the decl ref references a thunk from a natively foreign
   /// declaration to Swift calling convention.
   bool isForeignToNativeThunk() const;

include/swift/SIL/SILVTableVisitor.h

@@ -86,7 +86,24 @@ template <class T> class SILVTableVisitor {
   void maybeAddMethod(FuncDecl *fd) {
     assert(!fd->hasClangNode());
 
-    maybeAddEntry(SILDeclRef(fd, SILDeclRef::Kind::Func));
+    SILDeclRef constant(fd, SILDeclRef::Kind::Func);
+    maybeAddEntry(constant);
+
+    for (auto *diffAttr : fd->getAttrs().getAttributes<DifferentiableAttr>()) {
+      auto jvpConstant = constant.asAutoDiffDerivativeFunction(
+          AutoDiffDerivativeFunctionIdentifier::get(
+              AutoDiffDerivativeFunctionKind::JVP,
+              diffAttr->getParameterIndices(),
+              diffAttr->getDerivativeGenericSignature(), fd->getASTContext()));
+      maybeAddEntry(jvpConstant);
+
+      auto vjpConstant = constant.asAutoDiffDerivativeFunction(
+          AutoDiffDerivativeFunctionIdentifier::get(
+              AutoDiffDerivativeFunctionKind::VJP,
+              diffAttr->getParameterIndices(),
+              diffAttr->getDerivativeGenericSignature(), fd->getASTContext()));
+      maybeAddEntry(vjpConstant);
+    }
   }
 
   void maybeAddConstructor(ConstructorDecl *cd) {
@@ -96,7 +113,24 @@ template <class T> class SILVTableVisitor {
     // The initializing entry point for designated initializers is only
     // necessary for super.init chaining, which is sufficiently constrained
     // to never need dynamic dispatch.
-    maybeAddEntry(SILDeclRef(cd, SILDeclRef::Kind::Allocator));
+    SILDeclRef constant(cd, SILDeclRef::Kind::Allocator);
+    maybeAddEntry(constant);
+
+    for (auto *diffAttr : cd->getAttrs().getAttributes<DifferentiableAttr>()) {
+      auto jvpConstant = constant.asAutoDiffDerivativeFunction(
+          AutoDiffDerivativeFunctionIdentifier::get(
+              AutoDiffDerivativeFunctionKind::JVP,
+              diffAttr->getParameterIndices(),
+              diffAttr->getDerivativeGenericSignature(), cd->getASTContext()));
+      maybeAddEntry(jvpConstant);
+
+      auto vjpConstant = constant.asAutoDiffDerivativeFunction(
+          AutoDiffDerivativeFunctionIdentifier::get(
+              AutoDiffDerivativeFunctionKind::VJP,
+              diffAttr->getParameterIndices(),
+              diffAttr->getDerivativeGenericSignature(), cd->getASTContext()));
+      maybeAddEntry(vjpConstant);
+    }
   }
 
   void maybeAddAccessors(AbstractStorageDecl *asd) {

include/swift/SIL/SILWitnessVisitor.h

@@ -122,14 +122,19 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
 
   void visitAbstractStorageDecl(AbstractStorageDecl *sd) {
     sd->visitOpaqueAccessors([&](AccessorDecl *accessor) {
-      if (SILDeclRef::requiresNewWitnessTableEntry(accessor))
+      if (SILDeclRef::requiresNewWitnessTableEntry(accessor)) {
         asDerived().addMethod(SILDeclRef(accessor, SILDeclRef::Kind::Func));
+        addAutoDiffDerivativeMethodsIfRequired(accessor,
+                                               SILDeclRef::Kind::Func);
+      }
     });
   }
 
   void visitConstructorDecl(ConstructorDecl *cd) {
-    if (SILDeclRef::requiresNewWitnessTableEntry(cd))
+    if (SILDeclRef::requiresNewWitnessTableEntry(cd)) {
       asDerived().addMethod(SILDeclRef(cd, SILDeclRef::Kind::Allocator));
+      addAutoDiffDerivativeMethodsIfRequired(cd, SILDeclRef::Kind::Allocator);
+    }
   }
 
   void visitAccessorDecl(AccessorDecl *func) {
@@ -138,8 +143,10 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
 
   void visitFuncDecl(FuncDecl *func) {
     assert(!isa<AccessorDecl>(func));
-    if (SILDeclRef::requiresNewWitnessTableEntry(func))
+    if (SILDeclRef::requiresNewWitnessTableEntry(func)) {
       asDerived().addMethod(SILDeclRef(func, SILDeclRef::Kind::Func));
+      addAutoDiffDerivativeMethodsIfRequired(func, SILDeclRef::Kind::Func);
+    }
   }
 
   void visitMissingMemberDecl(MissingMemberDecl *placeholder) {
@@ -166,6 +173,26 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
   void visitPoundDiagnosticDecl(PoundDiagnosticDecl *pdd) {
     // We don't care about diagnostics at this stage.
   }
+
+private:
+  void addAutoDiffDerivativeMethodsIfRequired(AbstractFunctionDecl *AFD,
+                                              SILDeclRef::Kind kind) {
+    SILDeclRef declRef(AFD, kind);
+    for (auto *diffAttr : AFD->getAttrs().getAttributes<DifferentiableAttr>()) {
+      asDerived().addMethod(declRef.asAutoDiffDerivativeFunction(
+          AutoDiffDerivativeFunctionIdentifier::get(
+              AutoDiffDerivativeFunctionKind::JVP,
+              diffAttr->getParameterIndices(),
+              diffAttr->getDerivativeGenericSignature(),
+              AFD->getASTContext())));
+      asDerived().addMethod(declRef.asAutoDiffDerivativeFunction(
+          AutoDiffDerivativeFunctionIdentifier::get(
+              AutoDiffDerivativeFunctionKind::VJP,
+              diffAttr->getParameterIndices(),
+              diffAttr->getDerivativeGenericSignature(),
+              AFD->getASTContext())));
+    }
+  }
 };
 
 } // end namespace swift

lib/IRGen/GenDiffWitness.cpp

@@ -39,15 +39,28 @@ void IRGenModule::emitSILDifferentiabilityWitness(
   ConstantInitBuilder builder(*this);
   auto diffWitnessContents = builder.beginStruct();
 
+  // TODO(TF-1211): Uncomment assertions after upstreaming differentiation
+  // transform.
+  // The mandatory differentiation transform canonicalizes differentiability
+  // witnesses and ensures that JVPs/VJPs are populated.
+  /*
   assert(dw->getJVP() &&
          "Differentiability witness definition should have JVP");
   assert(dw->getVJP() &&
          "Differentiability witness definition should have VJP");
-
   diffWitnessContents.addBitCast(
       getAddrOfSILFunction(dw->getJVP(), NotForDefinition), Int8PtrTy);
   diffWitnessContents.addBitCast(
       getAddrOfSILFunction(dw->getVJP(), NotForDefinition), Int8PtrTy);
+  */
+  llvm::Constant *jvpValue = llvm::UndefValue::get(Int8PtrTy);
+  llvm::Constant *vjpValue = llvm::UndefValue::get(Int8PtrTy);
+  if (auto *jvpFn = dw->getJVP())
+    jvpValue = getAddrOfSILFunction(dw->getJVP(), NotForDefinition);
+  if (auto *vjpFn = dw->getJVP())
+    vjpValue = getAddrOfSILFunction(dw->getVJP(), NotForDefinition);
+  diffWitnessContents.addBitCast(jvpValue, Int8PtrTy);
+  diffWitnessContents.addBitCast(vjpValue, Int8PtrTy);
 
   getAddrOfDifferentiabilityWitness(
       dw, diffWitnessContents.finishAndCreateFuture());

lib/IRGen/GenKeyPath.cpp

@@ -992,8 +992,7 @@ emitKeyPathComponent(IRGenModule &IGM,
         auto methodProto = cast<ProtocolDecl>(dc);
         auto &protoInfo = IGM.getProtocolInfo(methodProto,
                                               ProtocolInfoKind::Full);
-        auto index = protoInfo.getFunctionIndex(
-                             cast<AbstractFunctionDecl>(declRef.getDecl()));
+        auto index = protoInfo.getFunctionIndex(declRef);
         idValue = llvm::ConstantInt::get(IGM.SizeTy, -index.getValue());
         idResolution = KeyPathComponentHeader::Resolved;
       }

lib/IRGen/GenProto.cpp

@@ -792,20 +792,19 @@ namespace {
     }
 
     void addMethod(SILDeclRef func) {
-      auto decl = cast<AbstractFunctionDecl>(func.getDecl());
       // If this assert needs to be changed, be sure to also change
       // ProtocolDescriptorBuilder::getRequirementInfo.
-      assert((isa<ConstructorDecl>(decl)
-                ? (func.kind == SILDeclRef::Kind::Allocator)
-                : (func.kind == SILDeclRef::Kind::Func))
-             && "unexpected kind for protocol witness declaration ref");
-      Entries.push_back(WitnessTableEntry::forFunction(decl));
+      assert((isa<ConstructorDecl>(func.getDecl())
+                  ? (func.kind == SILDeclRef::Kind::Allocator)
+                  : (func.kind == SILDeclRef::Kind::Func)) &&
+             "unexpected kind for protocol witness declaration ref");
+      Entries.push_back(WitnessTableEntry::forFunction(func));
     }
 
     void addPlaceholder(MissingMemberDecl *placeholder) {
       for (auto i : range(placeholder->getNumberOfVTableEntries())) {
         (void)i;
-        Entries.push_back(WitnessTableEntry());
+        Entries.push_back(WitnessTableEntry::forPlaceholder());
       }
     }
 
@@ -1318,8 +1317,7 @@ class AccessorConformanceInfo : public ConformanceInfo {
              && "sil witness table does not match protocol");
       assert(entry.getMethodWitness().Requirement == requirement
              && "sil witness table does not match protocol");
-      auto piIndex =
-        PI.getFunctionIndex(cast<AbstractFunctionDecl>(requirement.getDecl()));
+      auto piIndex = PI.getFunctionIndex(requirement);
       assert((size_t)piIndex.getValue() ==
               Table.size() - WitnessTableFirstRequirementOffset &&
              "offset doesn't match ProtocolInfo layout");
@@ -3277,7 +3275,7 @@ FunctionPointer irgen::emitWitnessMethodValue(IRGenFunction &IGF,
 
   // Find the witness we're interested in.
   auto &fnProtoInfo = IGF.IGM.getProtocolInfo(proto, ProtocolInfoKind::Full);
-  auto index = fnProtoInfo.getFunctionIndex(fn);
+  auto index = fnProtoInfo.getFunctionIndex(member);
   llvm::Value *slot;
   llvm::Value *witnessFnPtr =
     emitInvariantLoadOfOpaqueWitness(IGF, wtable,