[AutoDiff] Enable derivative registration for imported C functions. (#29016)

dan-zheng · web-flow · commit fe6998d8d436 · 2020-01-06T11:55:43.000-08:00
- Edit `SILFunctionType::getAutoDiffDerivativeFunctionType`. - Return `@convention(thin)` derivative function types given `@convention(c)` original function types. - Necessary because derivative function types have multiple results, which is unsupported for `@convention(c)` functions. - Revert derivative symbol TBDGen to visit `@differentiable` and `@derivative` attributes instead of visiting derivative configurations. - Necessary for `@derivative` attributes with cross-file original functions. `TBDGenVisitor` does not visit functions in other files, so the derivative configurations of the cross-file original function would not be visited. - SILGen: strip external from differentiability witness linkages. - Clang-imported functions have `public_external` linkage, but their differentiability witnesses should have `public` linkage. - Set `SILDeclRef` `isForeign` correctly for original functions that require a foreign entry point (including imported functions). - Done in SILGen, TBDGen, derivative lookup, IRGen. Resolves TF-1087: `@derivative` attribute SILGen crash for foreign functions. Unblocks TF-1085: using `@derivative` attribute for stdlib tgmath functions.
diff --git a/lib/IRGen/GenDecl.cpp b/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);
diff --git a/lib/IRGen/GenDiffWitness.cpp b/lib/IRGen/GenDiffWitness.cpp
@@ -34,9 +34,8 @@ void IRGenModule::emitSILDifferentiabilityWitness(
   if (dw->isDeclaration())
     return;
 
-  // Don't emit public_external witnesses.
-  if (hasPublicVisibility(dw->getLinkage()) &&
-      isAvailableExternally(dw->getLinkage()))
+  // Don't emit `public_external` witnesses.
+  if (dw->getLinkage() == SILLinkage::PublicExternal)
     return;
 
   ConstantInitBuilder builder(*this);
diff --git a/lib/SIL/SILFunctionType.cpp b/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 (getRepresentation() == SILFunctionTypeRepresentation::CFunctionPointer)
+    extInfo = extInfo.withRepresentation(SILFunctionTypeRepresentation::Thin);
+  return SILFunctionType::get(canGenSig, extInfo, getCoroutineKind(),
                               getCalleeConvention(), newParameters, getYields(),
                               newResults, getOptionalErrorResult(),
                               getSubstitutions(), isGenericSignatureImplied(),
diff --git a/lib/SILGen/SILGen.cpp b/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.
@@ -796,7 +793,9 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
           break;
         }
         auto *origAFD = derivAttr->getOriginalFunction();
-        auto *origFn = getFunction(SILDeclRef(origAFD), NotForDefinition);
+        auto origDeclRef =
+            SILDeclRef(origAFD).asForeign(requiresForeignEntryPoint(origAFD));
+        auto *origFn = getFunction(origDeclRef, NotForDefinition);
         auto derivativeGenSig = AFD->getGenericSignature();
         auto *resultIndices = IndexSubset::get(getASTContext(), 1, {0});
         AutoDiffConfig config(derivAttr->getParameterIndices(), resultIndices,
@@ -854,10 +853,13 @@ void SILGenModule::emitDifferentiabilityWitness(
   SILDifferentiabilityWitnessKey key{originalFunction->getName(), silConfig};
   auto *diffWitness = M.lookUpDifferentiabilityWitness(key);
   if (!diffWitness) {
+    // Strip external from linkage of original function.
+    // Necessary for Clang-imported functions, which have external linkage.
+    auto linkage = stripExternalFromLinkage(originalFunction->getLinkage());
     diffWitness = SILDifferentiabilityWitness::createDefinition(
-        M, originalFunction->getLinkage(), originalFunction,
-        silConfig.parameterIndices, silConfig.resultIndices,
-        config.derivativeGenericSignature, /*jvp*/ nullptr, /*vjp*/ nullptr,
+        M, linkage, originalFunction, silConfig.parameterIndices,
+        silConfig.resultIndices, config.derivativeGenericSignature,
+        /*jvp*/ nullptr, /*vjp*/ nullptr,
         /*isSerialized*/ hasPublicVisibility(originalFunction->getLinkage()),
         attr);
   }
@@ -881,8 +883,10 @@ void SILGenModule::emitDifferentiabilityWitness(
       auto *id = AutoDiffDerivativeFunctionIdentifier::get(
           kind, config.parameterIndices, config.derivativeGenericSignature,
           getASTContext());
+      auto origDeclRef = SILDeclRef(originalAFD)
+                             .asForeign(requiresForeignEntryPoint(originalAFD));
       derivativeThunk = getOrCreateAutoDiffDerivativeForwardingThunk(
-          SILDeclRef(originalAFD).asAutoDiffDerivativeFunction(id), derivative,
+          origDeclRef.asAutoDiffDerivativeFunction(id), derivative,
           expectedDerivativeType);
     }
     // Check for existing same derivative.
diff --git a/lib/SILOptimizer/Utils/Differentiation/Common.cpp b/lib/SILOptimizer/Utils/Differentiation/Common.cpp
@@ -331,7 +331,7 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
   if (resultIndices->getCapacity() != 1 || !resultIndices->contains(0))
     return nullptr;
 
-  // Explicit differentiability witnesses only exist on SILFunctions that come
+  // Explicit differentiability witnesses only exist on SIL functions that come
   // from AST functions.
   auto *originalAFD = findAbstractFunctionDecl(original);
   if (!originalAFD)
@@ -343,8 +343,16 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
   if (!minimalConfig)
     return nullptr;
 
-  auto *existingWitness = module.lookUpDifferentiabilityWitness(
-      {original->getName(), *minimalConfig});
+  std::string originalName = original->getName();
+  // If original function requires a foreign entry point, use the foreign SIL
+  // function to get or create the minimal differentiability witness.
+  if (requiresForeignEntryPoint(originalAFD)) {
+    originalName = SILDeclRef(originalAFD).asForeign().mangle();
+    original = module.lookUpFunction(SILDeclRef(originalAFD).asForeign());
+  }
+
+  auto *existingWitness =
+      module.lookUpDifferentiabilityWitness({originalName, *minimalConfig});
   if (existingWitness)
     return existingWitness;
 
@@ -358,6 +366,5 @@ SILDifferentiabilityWitness *getOrCreateMinimalASTDifferentiabilityWitness(
       minimalConfig->derivativeGenericSignature);
 }
 
-
 } // end namespace autodiff
 } // end namespace swift
diff --git a/lib/SILOptimizer/Utils/Differentiation/LinearMapInfo.cpp b/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;
 
diff --git a/lib/TBDGen/TBDGen.cpp b/lib/TBDGen/TBDGen.cpp
@@ -241,7 +241,8 @@ void TBDGenVisitor::addAutoDiffLinearMapFunction(AbstractFunctionDecl *original,
                                                  AutoDiffConfig config,
                                                  AutoDiffLinearMapKind kind) {
   auto &ctx = original->getASTContext();
-  auto declRef = SILDeclRef(original);
+  auto declRef =
+      SILDeclRef(original).asForeign(requiresForeignEntryPoint(original));
 
   if (!declRef.isSerialized())
     return;
@@ -270,31 +271,46 @@ void TBDGenVisitor::addAutoDiffDerivativeFunction(
   auto *assocFnId = AutoDiffDerivativeFunctionIdentifier::get(
       kind, parameterIndices, derivativeGenericSignature,
       original->getASTContext());
-  addSymbol(SILDeclRef(original).asAutoDiffDerivativeFunction(assocFnId));
+  auto declRef =
+      SILDeclRef(original).asForeign(requiresForeignEntryPoint(original));
+  addSymbol(declRef.asAutoDiffDerivativeFunction(assocFnId));
 }
 
 void TBDGenVisitor::addDifferentiabilityWitness(
     AbstractFunctionDecl *original, IndexSubset *astParameterIndices,
     IndexSubset *resultIndices, GenericSignature derivativeGenericSignature) {
-  if (SILDeclRef(original).getLinkage(ForDefinition) != SILLinkage::Public)
+  bool foreign = requiresForeignEntryPoint(original);
+  auto declRef = SILDeclRef(original).asForeign(foreign);
+
+  // Skip symbol emission for original functions that do not have public
+  // linkage. Exclude original functions that require a foreign entry point with
+  // `public_external` linkage.
+  auto originalLinkage = declRef.getLinkage(ForDefinition);
+  if (foreign)
+    originalLinkage = stripExternalFromLinkage(originalLinkage);
+  if (originalLinkage != SILLinkage::Public)
     return;
 
   auto *silParamIndices = autodiff::getLoweredParameterIndices(
       astParameterIndices,
       original->getInterfaceType()->castTo<AnyFunctionType>());
 
-  std::string originalMangledName = SILDeclRef(original).mangle();
+  auto originalMangledName = declRef.mangle();
   AutoDiffConfig config{silParamIndices, resultIndices,
                         derivativeGenericSignature};
   SILDifferentiabilityWitnessKey key(originalMangledName, config);
 
   Mangle::ASTMangler mangler;
-  std::string mangledName = mangler.mangleSILDifferentiabilityWitnessKey(key);
+  auto mangledName = mangler.mangleSILDifferentiabilityWitnessKey(key);
   addSymbol(mangledName);
 }
 
 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 +390,21 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
   }
 
   // SWIFT_ENABLE_TENSORFLOW
-  for (auto derivativeConfig : AFD->getDerivativeFunctionConfigurations()) {
-    addDerivativeConfiguration(AFD, derivativeConfig);
-  }
+  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()));
+  // SWIFT_ENABLE_TENSORFLOW END
 
   visitDefaultArguments(AFD, AFD->getParameters());
 }
@@ -430,6 +458,16 @@ 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()));
+  // SWIFT_ENABLE_TENSORFLOW END
 }
 
 void TBDGenVisitor::visitVarDecl(VarDecl *VD) {
diff --git a/lib/TBDGen/TBDGenVisitor.h b/lib/TBDGen/TBDGenVisitor.h
@@ -56,6 +56,15 @@ class TBDGenVisitor : public ASTVisitor<TBDGenVisitor> {
   const TBDGenOptions &Opts;
   Decl* TopLevelDecl = nullptr;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// A set of original function and derivative configuration pairs for which
+  /// derivative symbols have been emitted.
+  ///
+  /// Used to deduplicate derivative symbol emission for `@differentiable` and
+  /// `@derivative` attributes.
+  llvm::DenseSet<std::pair<AbstractFunctionDecl *, AutoDiffConfig>>
+      AddedDerivatives;
+
 private:
   void addSymbolInternal(StringRef name, llvm::MachO::SymbolKind kind,
                          bool isLinkerDirective = false);
diff --git a/test/AutoDiff/derivative_registration_foreign/Inputs/Foreign.c b/test/AutoDiff/derivative_registration_foreign/Inputs/Foreign.c
@@ -0,0 +1 @@
+float cFunction(float x) { return x; }
diff --git a/test/AutoDiff/derivative_registration_foreign/Inputs/Foreign.h b/test/AutoDiff/derivative_registration_foreign/Inputs/Foreign.h
@@ -0,0 +1 @@
+float cFunction(float);
diff --git a/test/AutoDiff/derivative_registration_foreign/Inputs/module.modulemap b/test/AutoDiff/derivative_registration_foreign/Inputs/module.modulemap
@@ -0,0 +1,4 @@
+module CForeign {
+  header "Foreign.h"
+  export *
+}
diff --git a/test/AutoDiff/derivative_registration_foreign/main.swift b/test/AutoDiff/derivative_registration_foreign/main.swift
@@ -0,0 +1,40 @@
+// RUN: %empty-directory(%t)
+// RUN: %clang -shared %S/Inputs/Foreign.c -fmodules -o %t/%target-library-name(CForeign)
+// RUN: %target-swift-emit-silgen -Xllvm -enable-experimental-cross-file-derivative-registration -I %S/Inputs -I %t %s | %FileCheck %s --check-prefix=CHECK-SILGEN --check-prefix=CHECK
+// RUN: %target-swift-emit-sil -Xllvm -enable-experimental-cross-file-derivative-registration -I %S/Inputs -I %t %s | %FileCheck %s --check-prefix=CHECK-SIL --check-prefix=CHECK
+// RUN: %target-build-swift -Xllvm -enable-experimental-cross-file-derivative-registration -I %S/Inputs -I %t %s -L %t -lCForeign
+
+import CForeign
+
+// TF-1087: Test derivative registration for foreign declaration (Clang-imported).
+// Original SILDeclRef must have `isForeign` bit set correctly.
+
+// CHECK-SILGEN-LABEL: // differentiability witness for cFunction
+// CHECK-SILGEN: sil_differentiability_witness [serialized] [parameters 0] [results 0] @cFunction : $@convention(c) (Float) -> Float {
+// CHECK-SILGEN:   vjp: @AD__cFunction__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-SILGEN: }
+
+// CHECK-SIL-LABEL: // differentiability witness for cFunction
+// CHECK-SIL: sil_differentiability_witness [serialized] [parameters 0] [results 0] @cFunction : $@convention(c) (Float) -> Float {
+// CHECK-SIL:   jvp: @AD__cFunction__jvp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-SIL:   vjp: @AD__cFunction__vjp_src_0_wrt_0 : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-SIL: }
+
+// Check that original SIL function is correct.
+
+// CHECK-SILGEN-LABEL: sil [serializable] [clang cFunction] @cFunction : $@convention(c) (Float) -> Float
+
+@inlinable
+@derivative(of: cFunction)
+func vjpCFunction(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (cFunction(x), { $0 })
+}
+
+@_silgen_name("test_derivative")
+@differentiable
+func testDerivative(_ x: Float) -> Float {
+  cFunction(x)
+}
+
+// CHECK-SILGEN-LABEL: sil hidden [ossa] @test_derivative : $@convention(thin) (Float) -> Float {
+// CHECK-SILGEN: {{%.*}} = function_ref @cFunction : $@convention(c) (Float) -> Float
