Introduce experimental features for autodiff options

Author: DougGregor

include/swift/Basic/Features.def

@@ -108,6 +108,14 @@ EXPERIMENTAL_FEATURE(TypeWitnessSystemInference)
 /// \endcode
 EXPERIMENTAL_FEATURE(BoundGenericExtensions)
 
+/// Whether to enable experimental differentiable programming features:
+/// `@differentiable` declaration attribute, etc.
+EXPERIMENTAL_FEATURE(DifferentiableProgramming)
+
+/// Whether to enable forward mode differentiation.
+EXPERIMENTAL_FEATURE(ForwardModeDifferentiation)
+
+
 #undef EXPERIMENTAL_FEATURE
 #undef FUTURE_FEATURE
 #undef SUPPRESSIBLE_LANGUAGE_FEATURE

include/swift/Basic/LangOptions.h

@@ -426,13 +426,6 @@ namespace swift {
     /// file.
     bool EmitFineGrainedDependencySourcefileDotFiles = false;
 
-    /// Whether to enable experimental differentiable programming features:
-    /// `@differentiable` declaration attribute, etc.
-    bool EnableExperimentalDifferentiableProgramming = false;
-
-    /// Whether to enable forward mode differentiation.
-    bool EnableExperimentalForwardModeDifferentiation = false;
-
     /// Whether to enable experimental `AdditiveArithmetic` derived
     /// conformances.
     bool EnableExperimentalAdditiveArithmeticDerivedConformances = false;

lib/AST/ASTPrinter.cpp

@@ -3039,6 +3039,14 @@ static bool usesFeatureBoundGenericExtensions(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureDifferentiableProgramming(Decl *decl) {
+  return false;
+}
+
+static bool usesFeatureForwardModeDifferentiation(Decl *decl) {
+  return false;
+}
+
 static void
 suppressingFeatureNoAsyncAvailability(PrintOptions &options,
                                       llvm::function_ref<void()> action) {

lib/AST/AutoDiff.cpp

@@ -112,7 +112,7 @@ void AutoDiffConfig::print(llvm::raw_ostream &s) const {
 bool swift::isDifferentiableProgrammingEnabled(SourceFile &SF) {
   auto &ctx = SF.getASTContext();
   // Return true if differentiable programming is explicitly enabled.
-  if (ctx.LangOpts.EnableExperimentalDifferentiableProgramming)
+  if (ctx.LangOpts.hasFeature(Feature::DifferentiableProgramming))
     return true;
   // Otherwise, return true iff the `_Differentiation` module is imported in
   // the given source file.

lib/Frontend/CompilerInvocation.cpp

@@ -547,9 +547,6 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_enable_experimental_additive_arithmetic_derivation))
     Opts.EnableExperimentalAdditiveArithmeticDerivedConformances = true;
 
-  Opts.EnableExperimentalForwardModeDifferentiation |=
-      Args.hasArg(OPT_enable_experimental_forward_mode_differentiation);
-
   Opts.DebuggerSupport |= Args.hasArg(OPT_debugger_support);
   if (Opts.DebuggerSupport)
     Opts.EnableDollarIdentifiers = true;
@@ -665,6 +662,8 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
     Opts.Features.insert(Feature::TypeWitnessSystemInference);
   if (Args.hasArg(OPT_enable_experimental_bound_generic_extensions))
     Opts.Features.insert(Feature::BoundGenericExtensions);
+  if (Args.hasArg(OPT_enable_experimental_forward_mode_differentiation))
+    Opts.Features.insert(Feature::ForwardModeDifferentiation);
 
   Opts.EnableAppExtensionRestrictions |= Args.hasArg(OPT_enable_app_extension);
 

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -898,7 +898,7 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     // - Functions with no return.
     // - Functions with unsupported control flow.
     if (context.getASTContext()
-            .LangOpts.EnableExperimentalForwardModeDifferentiation &&
+            .LangOpts.hasFeature(Feature::ForwardModeDifferentiation) &&
         (diagnoseNoReturn(context, witness->getOriginalFunction(), invoker) ||
          diagnoseUnsupportedControlFlow(
              context, witness->getOriginalFunction(), invoker)))
@@ -914,7 +914,7 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     // generation because generated JVP may not match semantics of custom VJP.
     // Instead, create an empty JVP.
     if (context.getASTContext()
-            .LangOpts.EnableExperimentalForwardModeDifferentiation &&
+            .LangOpts.hasFeature(Feature::ForwardModeDifferentiation) &&
         !witness->getVJP()) {
       // JVP and differential generation do not currently support functions with
       // multiple basic blocks.

lib/TBDGen/TBDGen.cpp

@@ -559,7 +559,7 @@ void TBDGenVisitor::addAutoDiffLinearMapFunction(AbstractFunctionDecl *original,
 
   // Differential functions are emitted only when forward-mode is enabled.
   if (kind == AutoDiffLinearMapKind::Differential &&
-      !ctx.LangOpts.EnableExperimentalForwardModeDifferentiation)
+      !ctx.LangOpts.hasFeature(Feature::ForwardModeDifferentiation))
     return;
   auto *loweredParamIndices = autodiff::getLoweredParameterIndices(
       config.parameterIndices,

tools/sil-opt/SILOpt.cpp

@@ -576,8 +576,10 @@ int main(int argc, char **argv) {
   if (EnableExperimentalStaticAssert)
     Invocation.getLangOptions().Features.insert(Feature::StaticAssert);
 
-  Invocation.getLangOptions().EnableExperimentalDifferentiableProgramming =
-      EnableExperimentalDifferentiableProgramming;
+  if (EnableExperimentalDifferentiableProgramming) {
+    Invocation.getLangOptions().Features.insert(
+        Feature::DifferentiableProgramming);
+  }
 
   Invocation.getLangOptions().EnableCXXInterop = EnableCxxInterop;