CMake: Support dynamic codegen outputs (pytorch#68246)

Summary:
Pull Request resolved: pytorch#68246

Currently the codegen produces a list of output files at CMake
configuration time and the build system has no way of knowing if the
outputs change. So if that happens, you basically need to delete the
build folder and re-run from scratch.

Instead, this generates the output list every time the code generation
is run and changes the output to be a `.cmake` file that gets included
in the main cmake configuration step. That means the build system
knows to re-run cmake automatically if a new output is added. So, for
example you could change the number of shards that `Operators.cpp` is
split into and it all just works transparently to the user.

Test Plan: Imported from OSS

Reviewed By: zou3519

Differential Revision: D32596268

Pulled By: albanD

fbshipit-source-id: 15e0896aeaead90aed64b9c8fda70cf28fef13a2

Author: peterbell10

cmake/Codegen.cmake

@@ -201,11 +201,31 @@ if(INTERN_BUILD_ATEN_OPS)
       ${GEN_VULKAN_FLAGS}
   )
 
+  file(GLOB_RECURSE headers_templates "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/templates/*\.h")
+  file(GLOB_RECURSE sources_templates "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/templates/*\.cpp")
+  set(declarations_yaml_templates "")
+
+  file(MAKE_DIRECTORY ${CMAKE_BINARY_DIR}/aten/src/ATen)
+  file(MAKE_DIRECTORY ${CMAKE_BINARY_DIR}/aten/src/ATen/core)
+
   foreach(gen_type "headers" "sources" "declarations_yaml")
+    # The codegen outputs may change dynamically as PyTorch is
+    # developed, but add_custom_command only supports dynamic inputs.
+    #
+    # We work around this by generating a .cmake file which is
+    # included below to set the list of output files. If that file
+    # ever changes then cmake will be re-run automatically because it
+    # was included and so we get fully dynamic outputs.
+
+    set("GEN_COMMAND_${gen_type}"
+        ${GEN_COMMAND}
+        --generate ${gen_type}
+        --output-dependencies ${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.cmake
+    )
+
+    # Dry run to bootstrap the output variables
     execute_process(
-        COMMAND ${GEN_COMMAND}
-          --generate ${gen_type}
-          --output-dependencies ${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.txt
+        COMMAND ${GEN_COMMAND_${gen_type}} --dry-run
         RESULT_VARIABLE RETURN_VALUE
         WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
     )
@@ -214,43 +234,22 @@ if(INTERN_BUILD_ATEN_OPS)
       message(FATAL_ERROR "Failed to get generated_${gen_type} list")
     endif()
 
-    file(READ "${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.txt" "generated_${gen_type}")
-    file(READ "${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.txt-cuda" "cuda_generated_${gen_type}")
-    file(READ "${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.txt-core" "core_generated_${gen_type}")
-  endforeach()
-
-  file(GLOB_RECURSE header_templates "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/templates/*\.h")
-  file(GLOB_RECURSE source_templates "${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/templates/*\.cpp")
-
-  file(MAKE_DIRECTORY ${CMAKE_BINARY_DIR}/aten/src/ATen)
-  file(MAKE_DIRECTORY ${CMAKE_BINARY_DIR}/aten/src/ATen/core)
-
-  add_custom_command(
-    OUTPUT ${generated_headers} ${cuda_generated_headers} ${core_generated_headers}
-    COMMAND ${GEN_COMMAND} --generate headers
-    DEPENDS ${all_python} ${header_templates}
-      ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/native_functions.yaml
-    WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
-    )
-
-  add_custom_command(
-    OUTPUT ${generated_sources} ${cuda_generated_sources} ${core_generated_sources}
-    COMMAND ${GEN_COMMAND} --generate sources
-    DEPENDS ${all_python} ${source_templates}
-      ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/native_functions.yaml
-    WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
-    )
-
-  add_custom_command(
-    OUTPUT
-      ${generated_declarations_yaml}
-      ${cuda_generated_declarations_yaml}
-      ${core_generated_declarations_yaml}
-    COMMAND ${GEN_COMMAND} --generate declarations_yaml
-    DEPENDS ${all_python}
-      ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/native_functions.yaml
-    WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
+    include("${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.cmake")
+    include("${CMAKE_BINARY_DIR}/aten/src/ATen/core_generated_${gen_type}.cmake")
+    include("${CMAKE_BINARY_DIR}/aten/src/ATen/cuda_generated_${gen_type}.cmake")
+
+    add_custom_command(
+      COMMENT "Generating ATen ${gen_type}"
+      OUTPUT ${generated_${gen_type}} ${cuda_generated_${gen_type}} ${core_generated_${gen_type}}
+        ${CMAKE_BINARY_DIR}/aten/src/ATen/generated_${gen_type}.cmake
+        ${CMAKE_BINARY_DIR}/aten/src/ATen/core_generated_${gen_type}.cmake
+        ${CMAKE_BINARY_DIR}/aten/src/ATen/cuda_generated_${gen_type}.cmake
+      COMMAND ${GEN_COMMAND_${gen_type}}
+      DEPENDS ${all_python} ${${gen_type}_templates}
+        ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/native_functions.yaml
+      WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
     )
+  endforeach()
 
   # Generated headers used from a CUDA (.cu) file are
   # not tracked correctly in CMake. We make the libATen.so depend explicitly

tools/codegen/gen.py

@@ -1200,6 +1200,9 @@ def main() -> None:
         '-o',
         '--output-dependencies',
         help='output a list of dependencies into the given file and exit')
+    parser.add_argument(
+        '--dry-run', action='store_true',
+        help='run without writing any files (still updates outputs)')
     parser.add_argument(
         '-d', '--install_dir', help='output directory',
         default='build/aten/src/ATen')
@@ -1271,7 +1274,7 @@ def main() -> None:
     pathlib.Path(core_install_dir).mkdir(parents=True, exist_ok=True)
 
     def make_file_manager(install_dir: str) -> FileManager:
-        return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.output_dependencies)
+        return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.dry_run)
 
     core_fm = make_file_manager(core_install_dir)
     cpu_fm = make_file_manager(options.install_dir)
@@ -1356,9 +1359,14 @@ def make_file_manager(install_dir: str) -> FileManager:
             cpu_fm=cpu_fm)
 
     if options.output_dependencies:
-        cpu_fm.write_outputs(options.output_dependencies)
-        core_fm.write_outputs(f"{options.output_dependencies}-core")
-        cuda_fm.write_outputs(f"{options.output_dependencies}-cuda")
+        depfile_path = pathlib.Path(options.output_dependencies).resolve()
+        depfile_name = depfile_path.name
+        depfile_stem = depfile_path.stem
+
+        for fm, prefix in [(cpu_fm, ""), (core_fm, "core_"), (cuda_fm, "cuda_")]:
+            varname = prefix + depfile_stem
+            path = depfile_path.parent / (prefix + depfile_name)
+            fm.write_outputs(varname, str(path))
 
 
 if __name__ == '__main__':

tools/codegen/utils.py

@@ -224,9 +224,9 @@ def merge_env(into: Dict[str, List[str]], from_: Dict[str, List[str]]) -> None:
         self.filenames.discard(
             f"{self.install_dir}/{base_filename}Everything{extension}")
 
-    def write_outputs(self, filename: str) -> None:
+    def write_outputs(self, variable_name: str, filename: str) -> None:
         """Write a file containing the list of all outputs which are
         generated by this script."""
-        self._write_if_changed(
-            filename,
-            ''.join(name + ";" for name in sorted(self.filenames)))
+        content = 'set({}\n    {})'.format(
+            variable_name, '\n    '.join('"' + name + '"' for name in sorted(self.filenames)))
+        self._write_if_changed(filename, content)