[WIP]

Try to make this pass:
```
import Darwin
print(gradient(at: Double(3), in: tan))
print(gradient(at: Double(3), in: Darwin.tan))
print(gradient(at: Double(3), in: Darwin.cos))
```

```
Undefined symbols for architecture x86_64:
  "_AD__$s6Darwin3cosyS2dF_PSRS", referenced from:
      _main in tgmath-e5045b.o
ld: symbol(s) not found for architecture x86_64
<unknown>:0: error: link command failed with exit code 1 (use -v to see invocation)
```

Author: dan-zheng

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 {

lib/IRGen/GenDecl.cpp

@@ -1078,9 +1078,6 @@ void IRGenerator::emitGlobalTopLevel(llvm::StringSet<> *linkerDirectives) {
   // Emit differentiability witnesses.
   for (auto &dw :
            PrimaryIGM->getSILModule().getDifferentiabilityWitnessList()) {
-    if (dw.isDeclaration())
-      continue;
-
     // Emit into same IRGenModule as the original function.
     // NOTE(TF-894): Investigate whether `getGenModule(dw.getVJP())` is
     // significant/desirable; `getGenModule` seems relevant for multi-threaded
@@ -4487,7 +4484,7 @@ IRGenModule::getAddrOfWitnessTablePattern(const NormalProtocolConformance *conf,
 }
 
 // SWIFT_ENABLE_TENSORFLOW
-/// Look up the address of a witness table.
+/// Look up the address of a differentiability witness.
 llvm::Constant *IRGenModule::getAddrOfDifferentiabilityWitness(
     const SILDifferentiabilityWitness *witness, ConstantInit definition) {
   auto entity = LinkEntity::forDifferentiabilityWitness(witness);

lib/IRGen/GenDiffWitness.cpp

@@ -29,15 +29,25 @@ void IRGenModule::emitSILDifferentiabilityWitness(
     SILDifferentiabilityWitness *dw) {
   PrettyStackTraceDifferentiabilityWitness _st(
       "emitting differentiability witness for", dw->getKey());
-
+  llvm::errs() << "IRGEN SIL DIFF WITNESS, CLANG DECL? "
+               << dw->getOriginalFunction()->getClangDecl() << ", C REFERENCES: "
+               << dw->getOriginalFunction()->hasCReferences() << ", WEAK IMPORTED: "
+               << dw->getOriginalFunction()->isWeakImported() << ", IS EXTERNAL DECL: "
+               << dw->getOriginalFunction()->isExternalDeclaration() << ", IS DECL: "
+               << dw->isDeclaration()
+               << "\n";
+  dw->dump();
   // Don't emit declarations.
   if (dw->isDeclaration())
     return;
 
-  // Don't emit public_external witnesses.
-  if (hasPublicVisibility(dw->getLinkage()) &&
-      isAvailableExternally(dw->getLinkage()))
-    return;
+  // Don't emit `public_external` witnesses.
+  // Make an exception for Clang-imported functions.
+  if (!dw->getOriginalFunction()->getClangDecl()) {
+    if (hasPublicVisibility(dw->getLinkage()) &&
+        isAvailableExternally(dw->getLinkage()))
+      return;
+  }
 
   ConstantInitBuilder builder(*this);
   auto diffWitnessContents = builder.beginStruct();
@@ -52,4 +62,6 @@ void IRGenModule::emitSILDifferentiabilityWitness(
 
   getAddrOfDifferentiabilityWitness(
       dw, diffWitnessContents.finishAndCreateFuture());
+
+  llvm::errs() << "IRGEN'd DIFF WITNESS!\n";
 }

lib/SIL/SILFunctionType.cpp

@@ -423,7 +423,13 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
   if (getSubstGenericSignature() && derivativeFnGenSig &&
       !derivativeFnGenSig->areAllParamsConcrete())
     canGenSig = derivativeFnGenSig;
-  return SILFunctionType::get(canGenSig, getExtInfo(), getCoroutineKind(),
+  // If original function is `@convention(c)`, the derivative function should
+  // have `@convention(thin)`. IRGen does not support `@convention(c)` functions
+  // with multiple results.
+  auto extInfo = getExtInfo();
+  if (getLanguage() == SILFunctionLanguage::C)
+    extInfo = extInfo.withRepresentation(SILFunctionTypeRepresentation::Thin);
+  return SILFunctionType::get(canGenSig, extInfo, getCoroutineKind(),
                               getCalleeConvention(), newParameters, getYields(),
                               newResults, getOptionalErrorResult(),
                               getSubstitutions(), isGenericSignatureImplied(),

lib/SILGen/SILGen.cpp

@@ -755,11 +755,8 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
              F->print(llvm::dbgs()));
 
   // SWIFT_ENABLE_TENSORFLOW
-  // Visit `@differentiable` attributes and generate SIL differentiability
-  // witnesses.
-  // TODO(TF-835): Visit `@derivative` attributes when type-checking no longer
-  // generates implicit `@differentiable` attributes. See TF-835 for replacement
-  // code.
+  // Visit `@differentiable` amd `@derivative` attributes and generate SIL
+  // differentiability witnesses.
   // Skip if the SILDeclRef is a:
   // - Default argument generator function.
   // - Thunk.
@@ -782,6 +779,10 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
                "all original SIL functions with generic signatures");
         AutoDiffConfig config(diffAttr->getParameterIndices(), resultIndices,
                               diffAttr->getDerivativeGenericSignature());
+        llvm::errs() << "SILGEN DIFF ATTR, ORIGINAL: " << AFD->getEffectiveFullName()
+        << ", F2N: " << constant.isForeignToNativeThunk()
+        << ", N2F: " << constant.isNativeToForeignThunk() << "\n";
+
         emitDifferentiabilityWitness(AFD, F, config, jvp, vjp, diffAttr);
       }
       for (auto *derivAttr : Attrs.getAttributes<DerivativeAttr>()) {
@@ -798,6 +799,15 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
         auto *origAFD = derivAttr->getOriginalFunction();
         auto origConstant =
             SILDeclRef(origAFD).asForeign(requiresForeignEntryPoint(origAFD));
+#if 0
+        auto origConstant = SILDeclRef(origAFD);
+#endif
+        llvm::errs() << "SILGEN DERIV ATTR, ORIGINAL: " << origAFD->getEffectiveFullName()
+        << ", F2N: " << origConstant.isForeignToNativeThunk()
+        << ", N2F: " << origConstant.isNativeToForeignThunk()
+        << ", HAS CLANG: " << origAFD->hasClangNode()
+        << ", FOREIGN: " << requiresForeignEntryPoint(origAFD)
+        << "\n";
         auto *origFn = getFunction(origConstant, NotForDefinition);
         auto derivativeGenSig = AFD->getGenericSignature();
         auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
@@ -863,6 +873,19 @@ void SILGenModule::emitDifferentiabilityWitness(
         /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
         attr);
   }
+  llvm::errs() << "SILGEN DIFF WITNESS, ORIG LINKAGE: "
+               << (unsigned)originalFunction->getLinkage() << "\n";
+  // Strip external from linkage of Clang-imported functions.
+  if (originalFunction->hasClangNode()) {
+    llvm::errs() << "SILGEN DIFF WITNESS, STRIP EXTERNAL FROM LINKAGE\n";
+    diffWitness->setLinkage(stripExternalFromLinkage(diffWitness->getLinkage()));
+  }
+  attr->print(llvm::errs(), originalAFD);
+  llvm::errs() << "\n";
+  diffWitness->dump();
+// #if 0
+  originalFunction->dump();
+// #endif
 
   // Set derivative function in differentiability witness.
   auto setDerivativeInDifferentiabilityWitness =

lib/SILOptimizer/Utils/Differentiation/Common.cpp

@@ -343,21 +343,32 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
   if (!minimalConfig)
     return nullptr;
 
+  std::string originalName = original->getName();
+  // If original function requires a foreign entry point, look up the
+  // differentiability witness for the foreign function.
+  if (requiresForeignEntryPoint(originalAFD)) {
+    originalName = SILDeclRef(originalAFD).asForeign().mangle();
+    original = module.lookUpFunction(SILDeclRef(originalAFD).asForeign());
+  }
+
   auto *existingWitness = module.lookUpDifferentiabilityWitness(
-      {original->getName(), *minimalConfig});
-  if (existingWitness)
+      {originalName, *minimalConfig});
+  if (existingWitness) {
+    llvm::errs() << "FOUND EXISTING WITNESS\n";
+    existingWitness->dump();
     return existingWitness;
+  }
 
   assert(original->isExternalDeclaration() &&
          "SILGen should create differentiability witnesses for all function "
          "definitions with explicit differentiable attributes");
+  llvm::errs() << "EXTERNAL WITNESS FOR: " << original->getName() << "\n";
 
   return SILDifferentiabilityWitness::createDeclaration(
       module, SILLinkage::PublicExternal, original,
       minimalConfig->parameterIndices, minimalConfig->resultIndices,
       minimalConfig->derivativeGenericSignature);
 }
 
-
 } // end namespace autodiff
 } // end namespace swift

lib/SILOptimizer/Utils/Differentiation/LinearMapInfo.cpp

@@ -472,7 +472,6 @@ bool LinearMapInfo::shouldDifferentiateApplySite(FullApplySite applySite) {
 
   // TODO: Pattern match to make sure there is at least one `store` to the
   // array's active buffer.
-  // if (isArrayLiteralIntrinsic(ai) && hasActiveResults)
   if (isArrayLiteralIntrinsic(applySite) && hasActiveResults)
     return true;
 

lib/TBDGen/TBDGen.cpp

@@ -65,6 +65,8 @@ void TBDGenVisitor::addSymbolInternal(StringRef name,
   if (StringSymbols && kind == SymbolKind::GlobalSymbol) {
     auto isNewValue = StringSymbols->insert(name).second;
     (void)isNewValue;
+    if (!isNewValue)
+      llvm::errs() << "DUPLICATE: " << name << "\n";
     assert(isNewValue && "symbol appears twice");
   }
 }
@@ -276,25 +278,55 @@ void TBDGenVisitor::addAutoDiffDerivativeFunction(
 void TBDGenVisitor::addDifferentiabilityWitness(
     AbstractFunctionDecl *original, IndexSubset *astParameterIndices,
     IndexSubset *resultIndices, GenericSignature derivativeGenericSignature) {
-  if (SILDeclRef(original).getLinkage(ForDefinition) != SILLinkage::Public)
+  bool foreign = requiresForeignEntryPoint(original);
+  llvm::errs() << "TBDGenVisitor::addDifferentiabilityWitness: "
+  << original->getEffectiveFullName() << ", HAS CLANG? " << original->hasClangNode() << ", FOREIGN: " << foreign << "\n";
+  if (!foreign) {
+  if (SILDeclRef(original).asForeign(foreign).getLinkage(ForDefinition) != SILLinkage::Public)
     return;
+  }
+  // original->dump();
+  
+#if 0
+  if (SILDeclRef(original).asForeign(requiresForeignEntryPoint(original)).getLinkage(ForDefinition) != SILLinkage::Public)
+    return;
+#endif
 
   auto *silParamIndices = autodiff::getLoweredParameterIndices(
       astParameterIndices,
       original->getInterfaceType()->castTo<AnyFunctionType>());
 
+  std::string originalMangledName = SILDeclRef(original).asForeign(requiresForeignEntryPoint(original)).mangle();
+#if 0
   std::string originalMangledName = SILDeclRef(original).mangle();
+#endif
   AutoDiffConfig config{silParamIndices, resultIndices,
                         derivativeGenericSignature};
   SILDifferentiabilityWitnessKey key(originalMangledName, config);
 
   Mangle::ASTMangler mangler;
   std::string mangledName = mangler.mangleSILDifferentiabilityWitnessKey(key);
   addSymbol(mangledName);
+#if 0
+  if (foreign) {
+    std::string originalMangledName = SILDeclRef(original).asForeign(true).mangle();
+    AutoDiffConfig config{silParamIndices, resultIndices,
+                          derivativeGenericSignature};
+    SILDifferentiabilityWitnessKey key(originalMangledName, config);
+
+    Mangle::ASTMangler mangler;
+    std::string mangledName = mangler.mangleSILDifferentiabilityWitnessKey(key);
+    addSymbol(mangledName);
+  }
+#endif
 }
 
 void TBDGenVisitor::addDerivativeConfiguration(AbstractFunctionDecl *original,
                                                AutoDiffConfig config) {
+  auto inserted = AddedDerivatives.insert({original, config});
+  if (!inserted.second)
+    return;
+
   addAutoDiffLinearMapFunction(original, config,
                                AutoDiffLinearMapKind::Differential);
   addAutoDiffLinearMapFunction(original, config,
@@ -374,9 +406,25 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
   }
 
   // SWIFT_ENABLE_TENSORFLOW
+#if 0
   for (auto derivativeConfig : AFD->getDerivativeFunctionConfigurations()) {
     addDerivativeConfiguration(AFD, derivativeConfig);
   }
+#endif
+  for (const auto *differentiableAttr :
+       AFD->getAttrs().getAttributes<DifferentiableAttr>())
+    addDerivativeConfiguration(
+        AFD,
+        AutoDiffConfig(differentiableAttr->getParameterIndices(),
+                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       differentiableAttr->getDerivativeGenericSignature()));
+  for (const auto *derivativeAttr :
+       AFD->getAttrs().getAttributes<DerivativeAttr>())
+    addDerivativeConfiguration(
+        derivativeAttr->getOriginalFunction(),
+        AutoDiffConfig(derivativeAttr->getParameterIndices(),
+                       IndexSubset::get(AFD->getASTContext(), 1, {0}),
+                       AFD->getGenericSignature()));
 
   visitDefaultArguments(AFD, AFD->getParameters());
 }
@@ -430,6 +478,15 @@ void TBDGenVisitor::visitAbstractStorageDecl(AbstractStorageDecl *ASD) {
   ASD->visitEmittedAccessors([&](AccessorDecl *accessor) {
     visitFuncDecl(accessor);
   });
+
+  // SWIFT_ENABLE_TENSORFLOW
+  for (const auto *differentiableAttr :
+       ASD->getAttrs().getAttributes<DifferentiableAttr>())
+    addDerivativeConfiguration(
+        ASD->getAccessor(AccessorKind::Get),
+        AutoDiffConfig(differentiableAttr->getParameterIndices(),
+                       IndexSubset::get(ASD->getASTContext(), 1, {0}),
+                       differentiableAttr->getDerivativeGenericSignature()));
 }
 
 void TBDGenVisitor::visitVarDecl(VarDecl *VD) {

lib/TBDGen/TBDGenVisitor.h

@@ -56,6 +56,15 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   const TBDGenOptions &Opts;
   Decl* TopLevelDecl = nullptr;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Tracks derivatives that have been added to the TBD.
+  ///
+  /// Different attributes trigger emission of the same derivatives (e.g.
+  /// `@differentiable` and `@derivative(of:)`), so we use this to deduplicate
+  /// the symbols associated with the derivatives in the TBD.
+  llvm::DenseSet<std::pair<AbstractFunctionDecl *, AutoDiffConfig>>
+      AddedDerivatives;
+
 private:
   void addSymbolInternal(StringRef name, llvm::MachO::SymbolKind kind,
                          bool isLinkerDirective = false);