[AutoDiff] flag for cross-file derivative registration

Author: Marc Rasi

lib/Sema/TypeCheckAttr.cpp

@@ -35,9 +35,16 @@
 #include "swift/Sema/IDETypeChecking.h"
 #include "clang/Basic/CharInfo.h"
 #include "llvm/Support/Debug.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "llvm/Support/Options.h"
 
 using namespace swift;
 
+// SWIFT_ENABLE_TENSORFLOW
+static llvm::cl::opt<bool> EnableExperimentalCrossFileDerivativeRegistration(
+    "enable-experimental-cross-file-derivative-registration",
+    llvm::cl::init(false));
+
 namespace {
   /// This emits a diagnostic with a fixit to remove the attribute.
   template<typename ...ArgTypes>
@@ -3646,7 +3653,8 @@ void AttributeChecker::visitDerivativeAttr(DerivativeAttr *attr) {
 
   // Reject different-file derivative registration.
   // TODO(TF-1021): Lift this restriction.
-  if (originalAFD->getParentSourceFile() != derivative->getParentSourceFile()) {
+  if (!EnableExperimentalCrossFileDerivativeRegistration &&
+      originalAFD->getParentSourceFile() != derivative->getParentSourceFile()) {
     diagnoseAndRemoveAttr(attr,
                           diag::derivative_attr_not_in_same_file_as_original);
     return;

test/AutoDiff/Inputs/cross_module_derivative_attr_e2e/main/main.swift

@@ -0,0 +1,12 @@
+import StdlibUnittest
+
+import module1
+
+var Tests = TestSuite("CrossModuleDerivativeAttr")
+
+Tests.test("CrossFile") {
+  let grad = gradient(at: 0, in: fCrossFile)
+  expectEqual(10, grad)
+}
+
+runAllTests()

test/AutoDiff/Inputs/cross_module_derivative_attr_e2e/module1/module1.swift

@@ -0,0 +1 @@
+public func fCrossFile(_ x: Float) -> Float { x }

test/AutoDiff/Inputs/cross_module_derivative_attr_e2e/module1/module1_other_file.swift

@@ -0,0 +1,4 @@
+@derivative(of: fCrossFile)
+public func vjpCrossFile(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  (x, { 10 * $0 })
+}

test/AutoDiff/cross_module_derivative_attr_e2e.swift

@@ -0,0 +1,5 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -I%t -c -parse-as-library -emit-module -module-name module1 -emit-module-path %t/module1.swiftmodule -o %t/module1.o %S/Inputs/cross_module_derivative_attr_e2e/module1/module1.swift %S/Inputs/cross_module_derivative_attr_e2e/module1/module1_other_file.swift -Xllvm -enable-experimental-cross-file-derivative-registration -validate-tbd-against-ir=none
+// RUN: %target-build-swift -I%t %S/Inputs/cross_module_derivative_attr_e2e/main/main.swift %t/module1.o -o %t/a.out -lm -Xllvm -enable-experimental-cross-file-derivative-registration -Xfrontend -validate-tbd-against-ir=none
+// RUN: %target-run %t/a.out
+// REQUIRES: executable_test