[AutoDiff] Add build-script flag for differentiable programming. (#27596)

Add `--enable-experimental-differentiable-programming` build-script flag.

The build-script flag enables/disables standard library additions
related to differentiable programming. This will allow official Swift
releases to disable these additions.

The build-script flag is on by default to ensure testing of
differentiable programming standard library additions. An additional
driver flag must be enabled to use differentiable programming features:
#27446

Mirror of #27595, with minor changes
specific to tensorflow branch.

Author: dan-zheng

CMakeLists.txt

@@ -392,6 +392,10 @@ option(SWIFT_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING
   "Build stdlibCore with exclusivity checking enabled"
   FALSE)
 
+option(SWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING
+  "Enable experimental Swift differentiable programming features"
+  FALSE)
+
 #
 # End of user-configurable options.
 #

cmake/modules/AddSwift.cmake

@@ -422,6 +422,10 @@ function(_add_variant_swift_compile_flags
     list(APPEND result "-D" "SWIFT_ENABLE_RUNTIME_FUNCTION_COUNTERS")
   endif()
 
+  if(SWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING)
+    list(APPEND result "-D" "SWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING")
+  endif()
+
   set("${result_var_name}" "${result}" PARENT_SCOPE)
 endfunction()
 

stdlib/public/core/CMakeLists.txt

@@ -197,6 +197,16 @@ set(SWIFTLIB_ESSENTIAL_GYB_SOURCES
   UnsafeRawBufferPointer.swift.gyb
   )
 
+# Compile differentiable programming sources only if enabled.
+set(SWIFTLIB_DIFFERENTIABLE_PROGRAMMING_SOURCES)
+if(SWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING)
+  list(APPEND SWIFTLIB_DIFFERENTIABLE_PROGRAMMING_SOURCES
+    Differentiable.swift
+    # SWIFT_ENABLE_TENSORFLOW
+    DifferentiationSupport.swift)
+  message(STATUS "Differentiable programming standard library additions enabled.")
+endif()
+
 # The complete list of sources in the core standard library.  Includes
 # all the essential sources listed above.
 set(SWIFTLIB_SOURCES
@@ -206,9 +216,6 @@ set(SWIFTLIB_SOURCES
   CollectionDifference.swift
   CollectionOfOne.swift
   Diffing.swift
-  Differentiable.swift
-  # SWIFT_ENABLE_TENSORFLOW
-  DifferentiationSupport.swift
   Mirror.swift
   PlaygroundDisplay.swift
   CommandLine.swift
@@ -218,6 +225,7 @@ set(SWIFTLIB_SOURCES
   VarArgs.swift
   Zip.swift
   "${SWIFT_SOURCE_DIR}/stdlib/linker-support/magic-symbols-for-install-name.c"
+  ${SWIFTLIB_DIFFERENTIABLE_PROGRAMMING_SOURCES}
   )
 
 set(SWIFTLIB_GYB_SOURCES

utils/build_swift/driver_arguments.py

@@ -1029,6 +1029,14 @@ def create_argument_parser():
                 'Currently only armv7 and aarch64 are supported. '
                 '%(default)s is the default.')
 
+    # -------------------------------------------------------------------------
+    in_group('Experimental language features')
+
+    option('--enable-experimental-differentiable-programming', toggle_true,
+           default=True,
+           help='Enable experimental Swift differentiable programming language'
+                ' features.')
+
     # -------------------------------------------------------------------------
     in_group('Unsupported options')
 

utils/build_swift/tests/expected_options.py

@@ -123,6 +123,7 @@
     'distcc': False,
     'dry_run': False,
     'enable_asan': False,
+    'enable_experimental_differentiable_programming': True,
     'enable_lsan': False,
     'enable_sanitize_coverage': False,
     'disable_guaranteed_normal_arguments': False,
@@ -459,6 +460,7 @@ class IgnoreOption(_BaseOption):
     EnableOption('--build-swift-stdlib-unittest-extra'),
     EnableOption('--distcc'),
     EnableOption('--enable-asan'),
+    EnableOption('--enable-experimental-differentiable-programming'),
     EnableOption('--enable-lsan'),
     EnableOption('--enable-sanitize-coverage'),
     EnableOption('--enable-tsan'),

utils/swift_build_support/swift_build_support/products/swift.py

@@ -41,6 +41,10 @@ def __init__(self, args, toolchain, source_dir, build_dir):
         # Add any exclusivity checking flags for stdlibcore.
         self.cmake_options.extend(self._stdlibcore_exclusivity_checking_flags)
 
+        # Add experimental differentiable programming flag.
+        self.cmake_options.extend(
+            self._enable_experimental_differentiable_programming)
+
     @property
     def _runtime_sanitizer_flags(self):
         sanitizer_list = []
@@ -112,3 +116,8 @@ def _force_optimized_typechecker_flags(self):
     def _stdlibcore_exclusivity_checking_flags(self):
         return [('SWIFT_STDLIB_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING:BOOL',
                  self.args.enable_stdlibcore_exclusivity_checking)]
+
+    @property
+    def _enable_experimental_differentiable_programming(self):
+        return [('SWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING:BOOL',
+                 self.args.enable_experimental_differentiable_programming)]

utils/swift_build_support/tests/products/test_swift.py

@@ -57,7 +57,8 @@ def setUp(self):
             benchmark_num_o_iterations=3,
             disable_guaranteed_normal_arguments=True,
             force_optimized_typechecker=False,
-            enable_stdlibcore_exclusivity_checking=False)
+            enable_stdlibcore_exclusivity_checking=False,
+            enable_experimental_differentiable_programming=False)
 
         # Setup shell
         shell.dry_run = True
@@ -87,7 +88,8 @@ def test_by_default_no_cmake_options(self):
         expected = [
             '-DCMAKE_EXPORT_COMPILE_COMMANDS=TRUE',
             '-DSWIFT_FORCE_OPTIMIZED_TYPECHECKER:BOOL=FALSE',
-            '-DSWIFT_STDLIB_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING:BOOL=FALSE'
+            '-DSWIFT_STDLIB_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING:BOOL=FALSE',
+            '-DSWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING:BOOL=FALSE'
         ]
         self.assertEqual(set(swift.cmake_options), set(expected))
 
@@ -102,7 +104,8 @@ def test_swift_runtime_tsan(self):
             '-DSWIFT_RUNTIME_USE_SANITIZERS=Thread',
             '-DCMAKE_EXPORT_COMPILE_COMMANDS=TRUE',
             '-DSWIFT_FORCE_OPTIMIZED_TYPECHECKER:BOOL=FALSE',
-            '-DSWIFT_STDLIB_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING:BOOL=FALSE'
+            '-DSWIFT_STDLIB_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING:BOOL=FALSE',
+            '-DSWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING:BOOL=FALSE'
         ]
         self.assertEqual(set(swift.cmake_options), set(flags_set))
 
@@ -302,3 +305,16 @@ def test_exclusivity_checking_flags(self):
              'TRUE'],
             [x for x in swift.cmake_options
              if 'SWIFT_STDLIB_ENABLE_STDLIBCORE_EXCLUSIVITY_CHECKING' in x])
+
+    def test_experimental_differentiable_programming_flags(self):
+        self.args.enable_experimental_differentiable_programming = True
+        swift = Swift(
+            args=self.args,
+            toolchain=self.toolchain,
+            source_dir='/path/to/src',
+            build_dir='/path/to/build')
+        self.assertEqual(
+            ['-DSWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING:BOOL='
+             'TRUE'],
+            [x for x in swift.cmake_options
+             if 'DSWIFT_ENABLE_EXPERIMENTAL_DIFFERENTIABLE_PROGRAMMING' in x])