[AutoDiff] Enable control flow AD by default. (#25273)

Preemptively remove `-Xllvm -differentiation-enable-control-flow` flag.

Note that limited control flow AD support is not yet robust:
- Memory leaks: TF-552.
- Incorrect zero gradients involving buffers + recursion: TF-554.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -62,14 +62,6 @@ static llvm::cl::opt<bool> SkipFoldingAutoDiffFunctionExtraction(
     "differentiation-skip-folding-autodiff-function-extraction",
     llvm::cl::init(false));
 
-/// This flag enables experimental differentiation of functions containing
-/// control flow. When this flag is true, adjoint generation is skipped and
-/// adjoint function references are replaced with undef. This is used for SIL
-/// testing purposes, specifically testing VJP generation before adjoint
-/// generation supports control flow.
-static llvm::cl::opt<bool> EnableControlFlow(
-    "differentiation-enable-control-flow", llvm::cl::init(false));
-
 //===----------------------------------------------------------------------===//
 // Helpers
 //===----------------------------------------------------------------------===//
@@ -1782,19 +1774,12 @@ static bool diagnoseUnsupportedControlFlow(ADContext &context,
         diag::autodiff_loops_not_supported);
     return true;
   }
-  // Diagnose unsupported terminators.
+  // Diagnose unsupported branching terminators.
   for (auto &bb : *original) {
     auto *term = bb.getTerminator();
-    // Adjoint generation is not yet robust for control flow.
-    // If control flow is not enable, emit an error.
-    // Otherwise, `br` and `cond_br` instructions are supported terminators.
-    if (isa<BranchInst>(term) || isa<CondBranchInst>(term)) {
-      if (EnableControlFlow)
-        continue;
-      context.emitNondifferentiabilityError(
-          term, invoker, diag::autodiff_control_flow_not_supported);
-      return true;
-    }
+    // Supported terminators are: `br`, `cond_br`.
+    if (isa<BranchInst>(term) || isa<CondBranchInst>(term))
+      continue;
     // If terminator is an unsupported branching terminator, emit an error.
     if (term->isBranch()) {
       context.emitNondifferentiabilityError(

test/AutoDiff/autodiff_diagnostics.swift

@@ -120,8 +120,6 @@ func calls_grad_of_nested(_ x: Float) -> Float {
 
 func if_else(_ x: Float, _ flag: Bool) -> Float {
   let y: Float
-  // expected-error @+2 {{expression is not differentiable}}
-  // expected-note @+1 {{differentiating control flow is not yet supported}}
   if flag {
     y = x + 1
   } else {

test/AutoDiff/control_flow.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift-control-flow-differentiation
+// RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
 import StdlibUnittest

test/AutoDiff/control_flow_diagnostics.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -differentiation-enable-control-flow %s
+// RUN: %target-swift-frontend -emit-sil -verify %s
 
 // Test supported `br` and `cond_br` terminators.
 

test/AutoDiff/control_flow_sil.swift

@@ -1,5 +1,5 @@
-// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -differentiation-enable-control-flow %s | %FileCheck %s -check-prefix=CHECK-DATA-STRUCTURES
-// RUN: %target-swift-frontend -emit-sil -verify -Xllvm -differentiation-enable-control-flow %s | %FileCheck %s -check-prefix=CHECK-SIL
+// RUN: %target-swift-frontend -emit-sil -verify %s | %FileCheck %s -check-prefix=CHECK-DATA-STRUCTURES
+// RUN: %target-swift-frontend -emit-sil -verify %s | %FileCheck %s -check-prefix=CHECK-SIL
 
 // TODO: Add adjoint SIL FileCheck tests.
 

test/AutoDiff/leakchecking.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift-control-flow-differentiation
+// RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
 // Test differentiation-related memory leaks.

test/TensorFlowRuntime/tensor_autodiff_control_flow.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift-control-flow-differentiation
+// RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 //
 // FIXME(TF-326): Re-enable `-O` after deserialization failure fix.

test/lit.cfg

@@ -1419,14 +1419,6 @@ if not getattr(config, 'target_run_simple_swift', None):
                             '%s %%t/a.out' % (config.target_build_swift,
                                               mcp_opt, config.target_codesign,
                                               config.target_run)))
-    # SWIFT_ENABLE_TENSORFLOW
-    # TODO: Remove when differentiation control flow support is robust.
-    config.target_run_simple_swift_control_flow_differentiation = (
-        '%%empty-directory(%%t) && '
-        '%s %s %%s -Xllvm -differentiation-enable-control-flow -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))
     config.target_run_simple_swift = (
         '%%empty-directory(%%t) && '
         '%s %s %%s -o %%t/a.out %s -module-name main  && '
@@ -1484,9 +1476,6 @@ config.substitutions.append(('%target-swift-frontend', config.target_swift_front
 
 
 config.substitutions.append(('%target-run-simple-swiftgyb', config.target_run_simple_swiftgyb))
-# SWIFT_ENABLE_TENSORFLOW
-# TODO: Remove when differentiation control flow support is robust.
-config.substitutions.append(('%target-run-simple-swift-control-flow-differentiation', config.target_run_simple_swift_control_flow_differentiation))
 config.substitutions.append(('%target-run-simple-swift\(([^)]+)\)', config.target_run_simple_swift_parameterized))
 config.substitutions.append(('%target-run-simple-swift', config.target_run_simple_swift))
 config.substitutions.append(('%target-run-stdlib-swiftgyb', config.target_run_stdlib_swiftgyb))