[AutoDiff] Change cross-file derivative registration to a frontend flag. (#29339)

Change `-enable-experimental-cross-file-derivative-registration` from an
`-Xllvm` flag to a `-Xfrontend` flag.

This makes the flag consistent with:
- `-enable-experimental-differentiable-programming`
- `-enable-experimental-forward-mode-differentiation`

Author: dan-zheng

cmake/modules/AddSwift.cmake

@@ -1914,7 +1914,7 @@ function(add_swift_target_library name)
     # SWIFT_ENABLE_TENSORFLOW
     # NOTE(TF-1021): Enable cross-file derivative registration for stdlib.
     list(APPEND swiftlib_swift_compile_flags_all
-         -Xllvm -enable-experimental-cross-file-derivative-registration)
+         -Xfrontend -enable-experimental-cross-file-derivative-registration)
     # SWIFT_ENABLE_TENSORFLOW END
 
     # Collect architecture agnostic SDK linker flags

include/swift/Basic/LangOptions.h

@@ -308,6 +308,9 @@ namespace swift {
     // SWIFT_ENABLE_TENSORFLOW END
 
     // SWIFT_ENABLE_TENSORFLOW
+    /// Whether to enable cross-file derivative registration.
+    bool EnableExperimentalCrossFileDerivativeRegistration = false;
+
     /// Whether to enable forward mode differentiation.
     bool EnableExperimentalForwardModeDifferentiation = false;
     // SWIFT_ENABLE_TENSORFLOW END

include/swift/Option/Options.td

@@ -470,9 +470,14 @@ def disable_bridging_pch : Flag<["-"], "disable-bridging-pch">,
 
 // Experimental feature options
 def enable_experimental_differentiable_programming : Flag<["-"], "enable-experimental-differentiable-programming">,
-  Flags<[FrontendOption]>,
+  Flags<[FrontendOption, ModuleInterfaceOption]>,
   HelpText<"Enable experimental differentiable programming features">;
+
 // SWIFT_ENABLE_TENSORFLOW
+def enable_experimental_cross_file_derivative_registration : Flag<["-"], "enable-experimental-cross-file-derivative-registration">,
+  Flags<[FrontendOption, ModuleInterfaceOption]>,
+  HelpText<"Enable experimental cross-file derivative registration">;
+
 // NOTE: This flag will be removed when JVP/differential generation is robust.
 def enable_experimental_forward_mode_differentiation : Flag<["-"], "enable-experimental-forward-mode-differentiation">,
   Flags<[FrontendOption]>,

lib/Frontend/CompilerInvocation.cpp

@@ -240,15 +240,6 @@ static void SaveModuleInterfaceArgs(ModuleInterfaceOptions &Opts,
     return;
   ArgStringList RenderedArgs;
   for (auto A : Args) {
-    // SWIFT_ENABLE_TENSORFLOW: Copy
-    // `-Xllvm -enable-experimental-cross-file-derivative-registration` into the
-    // .swiftinterface file.
-    if (A->getOption().matches(options::OPT_Xllvm) &&
-        StringRef(A->getValue()) ==
-            "-enable-experimental-cross-file-derivative-registration") {
-      A->render(Args, RenderedArgs);
-      continue;
-    }
     if (A->getOption().hasFlag(options::ModuleInterfaceOption))
       A->render(Args, RenderedArgs);
   }
@@ -382,9 +373,13 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_enable_experimental_differentiable_programming))
     Opts.EnableExperimentalDifferentiableProgramming = true;
 
-  // TODO: Ignore if enable-experimental-differentiable-programming is false
+  // TODO: Ignore differentiation-related flags if
+  // `enable-experimental-differentiable-programming` is false.
+  Opts.EnableExperimentalCrossFileDerivativeRegistration |=
+      Args.hasArg(OPT_enable_experimental_cross_file_derivative_registration);
   Opts.EnableExperimentalForwardModeDifferentiation |=
       Args.hasArg(OPT_enable_experimental_forward_mode_differentiation);
+
   if (Args.hasArg(OPT_enable_experimental_quasiquotes))
     Opts.EnableExperimentalQuasiquotes = true;
 

lib/Sema/TypeCheckAttr.cpp

@@ -3797,7 +3797,7 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
 
   // Reject different-file derivative registration.
   // TODO(TF-1021): Lift same-file derivative registration restriction.
-  if (!EnableExperimentalCrossFileDerivativeRegistration &&
+  if (!ctx.LangOpts.EnableExperimentalCrossFileDerivativeRegistration &&
       originalAFD->getParentSourceFile() != derivative->getParentSourceFile()) {
     diags.diagnose(attr->getLocation(),
                    diag::derivative_attr_not_in_same_file_as_original);

test/AutoDiff/downstream/cross_module_derivative_attr_e2e.swift

@@ -1,6 +1,6 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-build-swift -working-directory %t -I%t -parse-as-library -emit-module -module-name module1 -emit-module-path %t/module1.swiftmodule -emit-library -static %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 -Xfrontend -validate-tbd-against-ir=none
-// RUN: %target-build-swift -I%t -L%t %S/Inputs/cross_module_derivative_attr_e2e/main/main.swift -o %t/a.out -lm -lmodule1 -Xllvm -enable-experimental-cross-file-derivative-registration -Xfrontend -validate-tbd-against-ir=none
+// RUN: %target-build-swift -working-directory %t -I%t -parse-as-library -emit-module -module-name module1 -emit-module-path %t/module1.swiftmodule -emit-library -static %S/Inputs/cross_module_derivative_attr_e2e/module1/module1.swift %S/Inputs/cross_module_derivative_attr_e2e/module1/module1_other_file.swift -Xfrontend -enable-experimental-cross-file-derivative-registration -Xfrontend -validate-tbd-against-ir=none
+// RUN: %target-build-swift -I%t -L%t %S/Inputs/cross_module_derivative_attr_e2e/main/main.swift -o %t/a.out -lm -lmodule1 -Xfrontend -enable-experimental-cross-file-derivative-registration -Xfrontend -validate-tbd-against-ir=none
 // RUN: %target-run %t/a.out
 // REQUIRES: executable_test
 

test/AutoDiff/downstream/derivative_registration_foreign/main.swift

@@ -1,8 +1,8 @@
 // RUN: %empty-directory(%t)
 // RUN: %clang -c %S/Inputs/Foreign.c -fmodules -o %t/CForeign.o
-// 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 %t/CForeign.o
+// RUN: %target-swift-emit-silgen -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 -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 -Xfrontend -enable-experimental-cross-file-derivative-registration -I %S/Inputs -I %t %s %t/CForeign.o
 
 import CForeign