---

yaml
---
r: 340965
b: refs/heads/rxwei-patch-1
c: 58e5175
h: refs/heads/master
i:
  340963: cc112c8

Author: dan-zheng

[refs]

@@ -1015,7 +1015,7 @@ refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-08-18-a: b10b1fce14385faa6d44f6b933e95
 refs/heads/rdar-43033749-fix-batch-mode-no-diags-swift-5.0-branch: a14e64eaad30de89f0f5f0b2a782eed7ecdcb255
 refs/heads/revert-19006-error-bridging-integer-type: 8a9065a3696535305ea53fe9b71f91cbe6702019
 refs/heads/revert-19050-revert-19006-error-bridging-integer-type: ecf752d54b05dd0a20f510f0bfa54a3fec3bcaca
-refs/heads/rxwei-patch-1: 2e8ef909a086639bc6dcbac46325b4211845ae29
+refs/heads/rxwei-patch-1: 58e5175cf6af8e45d73d832d3c6a01840b40f07c
 refs/heads/shahmishal-patch-1: e58ec0f7488258d42bef51bc3e6d7b3dc74d7b2a
 refs/heads/typelist-existential: 4046359efd541fb5c72d69a92eefc0a784df8f5e
 refs/tags/swift-4.2-DEVELOPMENT-SNAPSHOT-2018-08-20-a: 4319ba09e4fb8650ee86061075c74a016b6baab9

branches/rxwei-patch-1/lib/SIL/SILFunctionBuilder.cpp

@@ -80,23 +80,42 @@ void SILFunctionBuilder::addFunctionAttributes(SILFunction *F,
       !constant.isStoredPropertyInitializer() &&
       !constant.isThunk()) {
     for (auto *A : Attrs.getAttributes<DifferentiableAttr>()) {
-      std::string jvpName, vjpName;
-      // Get JVP/VJP names.
-      if (auto *jvpFn = A->getJVPFunction())
-        jvpName = SILDeclRef(jvpFn).mangle();
-      if (auto *vjpFn = A->getVJPFunction())
-        vjpName = SILDeclRef(vjpFn).mangle();
       // Get lowered argument indices.
-      auto paramIndices = A->getParameterIndices();
+      auto *paramIndices = A->getParameterIndices();
       // NOTE: If `A->getParameterIndices()` is `nullptr`, continue. This is a
       // necessary hack regarding deserialization.
       if (!paramIndices)
         continue;
-      auto loweredParamIndices = paramIndices->getLowered(
+      auto *loweredParamIndices = paramIndices->getLowered(
           F->getASTContext(),
           decl->getInterfaceType()->castTo<AnyFunctionType>());
       SILAutoDiffIndices indices(/*source*/ 0, loweredParamIndices);
-      auto silDiffAttr = SILDifferentiableAttr::create(
+      // Get JVP/VJP names.
+      std::string jvpName, vjpName;
+      // If a method-self-reordering thunk is generated for the original
+      // function, use mangled JVP/VJP symbols.
+      auto *AFD = constant.getAbstractFunctionDecl();
+      auto selfParamIndex =
+          F->getLoweredFunctionType()->getNumParameters() - 1;
+      if (AFD && AFD->isInstanceMember() &&
+          F->getLoweredFunctionType()->hasSelfParam() &&
+          indices.isWrtParameter(selfParamIndex) &&
+          indices.parameters->getNumIndices() > 1) {
+        auto &ctx = F->getASTContext();
+        if (A->getJVPFunction())
+          jvpName = ctx.getIdentifier(
+              "AD__" + constant.mangle() + "__jvp_" + indices.mangle()).str();
+        if (A->getVJPFunction()) {
+          vjpName = ctx.getIdentifier(
+              "AD__" + constant.mangle() + "__vjp_" + indices.mangle()).str();
+        }
+      } else {
+        if (auto *jvpFn = A->getJVPFunction())
+          jvpName = SILDeclRef(jvpFn).mangle();
+        if (auto *vjpFn = A->getVJPFunction())
+          vjpName = SILDeclRef(vjpFn).mangle();
+      }
+      auto *silDiffAttr = SILDifferentiableAttr::create(
           M, indices, A->getRequirements(), M.allocateCopy(jvpName),
           M.allocateCopy(vjpName));
 #ifndef NDEBUG

branches/rxwei-patch-1/lib/SILGen/SILGen.cpp

@@ -781,6 +781,7 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
   LLVM_DEBUG(llvm::dbgs() << "lowered sil:\n";
              F->print(llvm::dbgs()));
 
+  // SWIFT_ENABLE_TENSORFLOW
   // Create self-reordering thunks for JVPs/VJPs of `@differentiable` methods.
   if (constant.hasDecl()) {
     auto *AFD = constant.getAbstractFunctionDecl();

branches/rxwei-patch-1/lib/TBDGen/TBDGen.cpp

@@ -219,15 +219,20 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
   // `@differentiable` attribute.
   auto diffAttrs = AFD->getAttrs().getAttributes<DifferentiableAttr>();
   for (auto *DA : diffAttrs) {
+    // If a method-self-reordering thunk is generated for the original function,
+    // emit symbol.
+    auto isSelfReorderedMethod = AFD && AFD->isInstanceMember() &&
+        AFD->hasImplicitSelfDecl() &&
+        DA->getParameterIndices()->parameters.count() > 1;
     // FIXME: When we get rid of `vjp:` and `jvp:` arguments in `@differentiable`,
     // we will no longer need to see whether they are specified.
-    if (!DA->getJVPFunction()) {
+    if (!DA->getJVPFunction() || isSelfReorderedMethod) {
       auto *id = AutoDiffAssociatedFunctionIdentifier::get(
           AutoDiffAssociatedFunctionKind::JVP, /*differentiationOrder*/ 1,
           DA->getParameterIndices(), AFD->getASTContext());
       addSymbol(SILDeclRef(AFD).asAutoDiffAssociatedFunction(id));
     }
-    if (!DA->getVJPFunction()) {
+    if (!DA->getVJPFunction() || isSelfReorderedMethod) {
       auto *id = AutoDiffAssociatedFunctionIdentifier::get(
           AutoDiffAssociatedFunctionKind::VJP, /*differentiationOrder*/ 1,
           DA->getParameterIndices(), AFD->getASTContext());
@@ -304,15 +309,21 @@ void TBDGenVisitor::visitAbstractStorageDecl(AbstractStorageDecl *ASD) {
   // with a `@differentiable` attribute.
   auto diffAttrs = ASD->getAttrs().getAttributes<DifferentiableAttr>();
   for (auto *DA : diffAttrs) {
+    // If a method-self-reordering thunk is generated for the original function,
+    // emit symbol.
+    auto isSelfReorderedMethod = ASD->getGetter() &&
+        ASD->getGetter()->isInstanceMember() &&
+        ASD->getGetter()->hasImplicitSelfDecl() &&
+        DA->getParameterIndices()->parameters.count() > 1;
     // FIXME: When we get rid of `vjp:` and `jvp:` arguments in `@differentiable`,
     // we will no longer need to see whether they are specified.
-    if (!DA->getJVPFunction()) {
+    if (!DA->getJVPFunction() || isSelfReorderedMethod) {
       auto *id = AutoDiffAssociatedFunctionIdentifier::get(
           AutoDiffAssociatedFunctionKind::JVP, /*differentiationOrder*/ 1,
           DA->getParameterIndices(), ASD->getASTContext());
       addSymbol(SILDeclRef(ASD->getGetter()).asAutoDiffAssociatedFunction(id));
     }
-    if (!DA->getVJPFunction()) {
+    if (!DA->getVJPFunction() || isSelfReorderedMethod) {
       auto *id = AutoDiffAssociatedFunctionIdentifier::get(
           AutoDiffAssociatedFunctionKind::VJP, /*differentiationOrder*/ 1,
           DA->getParameterIndices(), ASD->getASTContext());

branches/rxwei-patch-1/test/AutoDiff/Inputs/method_self_reordering_thunk_other_module.swift

@@ -0,0 +1,12 @@
+struct TF_619: Differentiable {
+  var p: Float = 1
+
+  @differentiable(vjp: vjpFoo)
+  func foo(_ x: Float) -> Float {
+    return p * x
+  }
+
+  func vjpFoo(_ x: Float) -> (Float, (Float) -> (TangentVector, Float)) {
+    return (x, { v in (TangentVector(p: v * x), v * self.p) })
+  }
+}

branches/rxwei-patch-1/test/AutoDiff/method_self_reordering_thunk/main.swift

@@ -0,0 +1,18 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %S/../Inputs/method_self_reordering_thunk_other_module.swift %s -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out
+
+// REQUIRES: executable_test
+
+import StdlibUnittest
+
+var MethodSelfReorderingThunkTests = TestSuite("MethodSelfReorderingThunks")
+
+// Test TF-619: cross-module import of `@differentiable` methods with
+// self-ordering thunks.
+MethodSelfReorderingThunkTests.test("CrossModuleMethodSelfReorderingThunk") {
+  expectEqual(1, gradient(at: 0) { x in TF_619().foo(x) })
+}
+
+runAllTests()

branches/rxwei-patch-1/test/AutoDiff/tbdgen.swift

@@ -42,6 +42,14 @@ public extension Float {
   }
 
   // This should generate public symbols for both JVP and VJP.
+  // Tests self-reordering-method thunking.
+  @differentiable
+  func method(x: Float, y: Float) -> Float {
+    return x
+  }
+
+  // This should generate public symbols for both JVP and VJP.
+  // Tests self-reordering-method thunking.
   @differentiable
   subscript(x: Float) -> Float {
     return x