[TF-935] Replaces '-Xllvm -run-jvp-generation' by '-enable-experiment… (#27861)

…al-forward-mode-differentiation'.

-enable-experimental-forward-mode-differentiation is a frontend option that is clearer on it's intent and more
accessible. It defaults to false.

Resolves [TF-935](https://bugs.swift.org/browse/TF-935).

Author: vguerra

include/swift/Basic/LangOptions.h

@@ -331,6 +331,11 @@ namespace swift {
     /// `@differentiable` declaration attribute, etc.
     bool EnableExperimentalDifferentiableProgramming = false;
 
+    // SWIFT_ENABLE_TENSORFLOW
+    /// Whether to enable forward mode differentiation.
+    bool EnableExperimentalForwardModeDifferentiation = false;
+    // SWIFT_ENABLE_TENSORFLOW END
+
     /// Whether to enable #quote, #unquote and @quoted.
     bool EnableExperimentalQuasiquotes = false;
 

include/swift/Option/Options.td

@@ -428,6 +428,12 @@ def disable_bridging_pch : Flag<["-"], "disable-bridging-pch">,
 def enable_experimental_differentiable_programming : Flag<["-"], "enable-experimental-differentiable-programming">,
   Flags<[FrontendOption]>,
   HelpText<"Enable experimental differentiable programming features">;
+// SWIFT_ENABLE_TENSORFLOW
+// 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]>,
+  HelpText<"Enable experimental forward mode differentiation">;
+// SWIFT_ENABLE_TENSORFLOW END
 
 def enable_experimental_quasiquotes : Flag<["-"], "enable-experimental-quasiquotes">,
   Flags<[FrontendOption]>,

lib/Driver/ToolChains.cpp

@@ -231,6 +231,10 @@ static void addCommonFrontendArgs(const ToolChain &TC, const OutputInfo &OI,
                        options::OPT_enable_experimental_dependencies);
   inputArgs.AddLastArg(arguments,
                        options::OPT_experimental_dependency_include_intrafile);
+  // SWIFT_ENABLE_TENSORFLOW
+  inputArgs.AddLastArg(
+      arguments, options::OPT_enable_experimental_forward_mode_differentiation);
+  // SWIFT_ENABLE_TENSORFLOW END
   inputArgs.AddLastArg(arguments,
                        options::OPT_enable_experimental_quasiquotes);
   inputArgs.AddLastArg(arguments, options::OPT_package_description_version);

lib/Frontend/CompilerInvocation.cpp

@@ -358,6 +358,9 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_experimental_dependency_include_intrafile))
     Opts.ExperimentalDependenciesIncludeIntrafileOnes = true;
 
+  // TODO: Ignore if enable-experimental-differentiable-programming is false
+  Opts.EnableExperimentalForwardModeDifferentiation |=
+      Args.hasArg(OPT_enable_experimental_forward_mode_differentiation);
   if (Args.hasArg(OPT_enable_experimental_quasiquotes))
     Opts.EnableExperimentalQuasiquotes = true;
 

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -64,12 +64,6 @@ static llvm::cl::opt<bool> SkipFoldingDifferentiableFunctionExtraction(
     "differentiation-skip-folding-differentiable-function-extraction",
     llvm::cl::init(true));
 
-/// This flag is used to enable full JVP generation.
-/// It will be removed when JVP/differential generation is robust.
-static llvm::cl::opt<bool> RunJVPGeneration(
-    "run-jvp-generation",
-    llvm::cl::init(false));
-
 //===----------------------------------------------------------------------===//
 // Helpers
 //===----------------------------------------------------------------------===//
@@ -8030,8 +8024,9 @@ bool ADContext::processDifferentiableAttribute(
     // Diagnose:
     // - Functions with no return.
     // - Functions with unsupported control flow.
-    if (RunJVPGeneration && (diagnoseNoReturn(*this, original, invoker) ||
-        diagnoseUnsupportedControlFlow(*this, original, invoker)))
+    if (getASTContext().LangOpts.EnableExperimentalForwardModeDifferentiation &&
+        (diagnoseNoReturn(*this, original, invoker) ||
+         diagnoseUnsupportedControlFlow(*this, original, invoker)))
       return true;
 
     jvp = createEmptyJVP(*this, original, attr, isDerivativeFnExported);
@@ -8041,7 +8036,8 @@ bool ADContext::processDifferentiableAttribute(
     // does not exist. If custom VJP exists but custom JVP does not, skip JVP
     // generation because generated JVP may not match semantics of custom VJP.
     // Instead, create an empty JVP.
-    if (RunJVPGeneration && !vjp) {
+    if (getASTContext().LangOpts.EnableExperimentalForwardModeDifferentiation &&
+        !vjp) {
       // JVP and differential generation do not currently support functions with
       // multiple basic blocks.
       if (original->getBlocks().size() > 1) {

test/AutoDiff/forward_mode_diagnostics.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -Xllvm -run-jvp-generation -emit-sil -verify %s
+// RUN: %target-swift-frontend -enable-experimental-forward-mode-differentiation -emit-sil -verify %s
 
 // TODO: move these tests back into `autodiff_diagnostics.swift` once
 // forward mode reaches feature parity with reverse mode.

test/AutoDiff/forward_mode_sil.swift

@@ -1,5 +1,5 @@
-// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -run-jvp-generation -Xllvm -debug-only=differentiation %s 2>&1 | %FileCheck %s -check-prefix=CHECK-DATA-STRUCTURES
-// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -sil-print-after=differentiation -Xllvm -run-jvp-generation -o /dev/null 2>&1 %s | %FileCheck %s -check-prefix=CHECK-SIL
+// RUN: %target-swift-frontend -emit-sil -verify -enable-experimental-forward-mode-differentiation -Xllvm -debug-only=differentiation %s 2>&1 | %FileCheck %s -check-prefix=CHECK-DATA-STRUCTURES
+// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -sil-print-after=differentiation -enable-experimental-forward-mode-differentiation -o /dev/null 2>&1 %s | %FileCheck %s -check-prefix=CHECK-SIL
 // REQUIRES: asserts
 
 

test/lit.cfg

@@ -1538,7 +1538,7 @@ if not getattr(config, 'target_run_simple_swift', None):
     # TODO: Remove when forward mode AD support is robust.	
     config.target_run_simple_swift_forward_mode_differentiation = (	
         '%%empty-directory(%%t) && '	
-        '%s %s %%s -Xllvm -run-jvp-generation -o %%t/a.out %s -module-name main  && '	
+        '%s %s %%s -enable-experimental-forward-mode-differentiation -o %%t/a.out %s -module-name main  && '
         '%s %%t/a.out &&'	
         '%s %%t/a.out'	
         % (config.target_build_swift, mcp_opt, swift_tensorflow_extra_options, config.target_codesign, config.target_run))