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
walterddr · Dec 7, 2021 · 9a7732e · 9a7732e
1 parent cd9da32
commit 9a7732e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 47 deletions.
diff --git a/cmake/Codegen.cmake b/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

diff --git a/tools/codegen/gen.py b/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__':

diff --git a/tools/codegen/utils.py b/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)