Update on "Move enum_tag to headeronly"

Differential Revision: [D86831013](https://our.internmc.facebook.com/intern/diff/D86831013)

[ghstack-poisoned]

.lintrunner.toml

@@ -243,6 +243,7 @@ exclude_patterns = [
     'torch/csrc/utils/generated_serialization_types.h',
     'torch/csrc/utils/pythoncapi_compat.h',
     'torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h',
+    "torch/headeronly/templates/*",
 ]
 init_command = [
     'python3',

BUILD.bazel

@@ -14,7 +14,10 @@ load(":ufunc_defs.bzl", "aten_ufunc_generated_cpu_kernel_sources", "aten_ufunc_g
 load("//:tools/bazel.bzl", "rules")
 
 # Export files for use by torch/headeronly (where version.h generation now lives)
-exports_files(["version.txt"])
+exports_files([
+    "version.txt",
+    "aten/src/ATen/native/tags.yaml",
+])
 
 define_targets(rules = rules)
 

aten/src/ATen/templates/enum_tag.h

@@ -1,10 +1,5 @@
 #pragma once
 
 // ${generated_comment}
-
-namespace at {
-    // Enum of valid tags obtained from the entries in tags.yaml
-    enum class Tag {
-        ${enum_of_valid_tags}
-    };
-}
+// This header forwards to the actual implementation in torch/headeronly
+#include <torch/headeronly/core/enum_tag.h>

buckbuild.bzl

@@ -1216,6 +1216,9 @@ def define_buck_targets(
             "core/enum_tag.h": ":gen_aten[core/enum_tag.h]",
         }),
         labels = labels,
+        exported_deps = [
+            "{}torch/headeronly:enum_tag_h".format(ROOT_PATH),
+        ],
     )
 
     fb_xplat_cxx_library(

build.bzl

@@ -65,6 +65,7 @@ def define_targets(rules):
     gen_aten_srcs = [
         "aten/src/ATen/native/native_functions.yaml",
         "aten/src/ATen/native/tags.yaml",
+        "//torch/headeronly:enum_tag_h",
     ] + rules.glob(["aten/src/ATen/templates/*"])
 
     gen_aten_cmd = " ".join([

cmake/Codegen.cmake

@@ -304,17 +304,52 @@ if(INTERN_BUILD_ATEN_OPS)
       DEPENDS ${all_python} ${${gen_type}_templates}
         ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/native_functions.yaml
         ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/tags.yaml
+        ${CMAKE_BINARY_DIR}/torch/headeronly/core/enum_tag.h
       WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
     )
   endforeach()
 
+  # Generate headeronly files (enum_tag.h)
+  set(HEADERONLY_BASE_DIR ${CMAKE_BINARY_DIR}/torch/headeronly)
+  set(GEN_HEADERONLY_COMMAND
+      "${Python_EXECUTABLE}" -m torchgen.gen
+      --source-path ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen
+      --headeronly
+      --headeronly-install-dir ${HEADERONLY_BASE_DIR}
+      --output-dependencies ${CMAKE_BINARY_DIR}/torch/headeronly/core/generated_headeronly.cmake
+  )
+
+  # Dry run to bootstrap the output variables
+  execute_process(
+      COMMAND ${GEN_HEADERONLY_COMMAND} --dry-run
+      RESULT_VARIABLE RETURN_VALUE
+      WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
+  )
+
+  if(NOT RETURN_VALUE EQUAL 0)
+    message(FATAL_ERROR "Failed to get generated headeronly list")
+  endif()
+
+  include("${CMAKE_BINARY_DIR}/torch/headeronly/core/generated_headeronly.cmake")
+
+  add_custom_command(
+    COMMENT "Generating headeronly files"
+    OUTPUT ${headeronly_generated_headers} ${CMAKE_BINARY_DIR}/torch/headeronly/core/generated_headeronly.cmake
+    COMMAND ${GEN_HEADERONLY_COMMAND}
+    DEPENDS ${all_python}
+      ${CMAKE_CURRENT_LIST_DIR}/../torch/headeronly/templates/enum_tag.h
+      ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/tags.yaml
+    WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
+  )
+
   # Generated headers used from a CUDA (.cu) file are
   # not tracked correctly in CMake. We make the libATen.so depend explicitly
   # on building the generated ATen files to workaround.
   add_custom_target(ATEN_CPU_FILES_GEN_TARGET DEPENDS
       ${generated_headers} ${core_generated_headers} ${cpu_vec_generated_headers} ${ops_generated_headers}
       ${generated_sources} ${core_generated_sources} ${cpu_vec_generated_sources} ${ops_generated_sources}
-      ${generated_declarations_yaml} ${generated_unboxing_sources})
+      ${generated_declarations_yaml} ${generated_unboxing_sources}
+      ${headeronly_generated_headers})
   add_custom_target(ATEN_CUDA_FILES_GEN_TARGET DEPENDS
       ${cuda_generated_headers} ${cuda_generated_sources})
   add_library(ATEN_CPU_FILES_GEN_LIB INTERFACE)

torch/headeronly/CMakeLists.txt

@@ -31,6 +31,9 @@ add_library(headeronly INTERFACE ${HEADERONLY_HEADERS})
 install(FILES ${CMAKE_BINARY_DIR}/torch/headeronly/macros/cmake_macros.h
         DESTINATION include/torch/headeronly/macros)
 
+install(FILES ${CMAKE_BINARY_DIR}/torch/headeronly/core/enum_tag.h
+        DESTINATION include/torch/headeronly/core)
+
 if(NOT BUILD_LIBTORCHLESS)
   # ---[ Installation copied from c10/CMakeLists.txt
   install(TARGETS headeronly EXPORT Caffe2Targets DESTINATION lib)

torch/headeronly/build.bzl

@@ -25,11 +25,32 @@ def define_targets(rules):
 
     rules.genrule(**genrule_args)
 
+    # Generate enum_tag.h using torchgen
+    enum_tag_genrule_args = {
+        "name": "enum_tag_h",
+        "srcs": [
+            "//aten/src/ATen/native:native_functions.yaml",
+            "//aten/src/ATen/native:tags.yaml",
+            "templates/enum_tag.h",
+        ],
+        "outs": ["core/enum_tag.h"],
+        "cmd": "$(execpath //torchgen:gen) " +
+               "--source-path aten/src/ATen " +
+               "--headeronly " +
+               "--headeronly-install-dir $(@D)/..",
+        "tools": ["//torchgen:gen"],
+    }
+
+    if not is_buck:
+        enum_tag_genrule_args["visibility"] = ["//visibility:public"]
+
+    rules.genrule(**enum_tag_genrule_args)
+
     rules.cc_library(
         name = "torch_headeronly",
         hdrs = rules.glob([
             "**/*.h"
-        ]) + ["version.h.in"],
+        ]) + ["version.h.in", "core/enum_tag.h"],
         visibility = ["//visibility:public"],
         deps = [
             "//torch/headeronly/macros",

torch/headeronly/templates/enum_tag.h

@@ -0,0 +1,14 @@
+#pragma once
+
+// ${generated_comment}
+
+#include <torch/headeronly/macros/Macros.h>
+
+namespace at {
+// Enum of valid tags obtained from the entries in tags.yaml
+enum class Tag { ${enum_of_valid_tags} };
+} // namespace at
+
+HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
+using at::Tag;
+HIDDEN_NAMESPACE_END(torch, headeronly)

torchgen/gen.py

@@ -2191,10 +2191,8 @@ def gen_headers(
 
     core_fm.write("aten_interned_strings.h", gen_aten_interned_strings)
 
-    def gen_tags_enum() -> dict[str, str]:
-        return {"enum_of_valid_tags": (",\n".join(sorted(valid_tags)))}
-
-    core_fm.write("enum_tag.h", gen_tags_enum)
+    # Generate forwarding header at ATen/core/enum_tag.h for backward compatibility
+    core_fm.write("enum_tag.h", dict)
 
 
 def gen_source_files(
@@ -2823,6 +2821,17 @@ def main() -> None:
         default=["headers", "sources", "declarations_yaml"],
         help="Generate only a subset of files",
     )
+    parser.add_argument(
+        "--headeronly",
+        action="store_true",
+        help="Generate only torch/headeronly generated files (e.g. enum_tag.h).",
+    )
+    parser.add_argument(
+        "--headeronly-install-dir",
+        "--headeronly_install_dir",
+        help="base output directory for headeronly files (core/ and templates/ will be derived from this)",
+        default="torch/headeronly",
+    )
     parser.add_argument(
         "--update-aoti-c-shim",
         action="store_true",
@@ -2966,6 +2975,37 @@ def main() -> None:
             if dp_key not in functions_keys:
                 functions_keys.add(dp_key)
 
+    if options.headeronly:
+        # Generate only the headeronly files (enum_tag.h)
+        # Parse tags.yaml to get valid tags
+        headeronly_valid_tags = parse_tags_yaml(tags_yaml_path)
+
+        # Create a FileManager for headeronly output
+        # The install directory is derived from the base, but the template directory
+        # is always in the source tree at torch/headeronly/templates
+        headeronly_base_dir = options.headeronly_install_dir
+        headeronly_install_dir = f"{headeronly_base_dir}/core"
+        Path(headeronly_install_dir).mkdir(parents=True, exist_ok=True)
+        headeronly_template_dir = "torch/headeronly/templates"
+        headeronly_fm = make_file_manager(
+            options=options,
+            install_dir=headeronly_install_dir,
+            template_dir=headeronly_template_dir,
+        )
+
+        # Generate enum_tag.h
+        def gen_tags_enum() -> dict[str, str]:
+            return {"enum_of_valid_tags": (",\n".join(sorted(headeronly_valid_tags)))}
+
+        headeronly_fm.write("enum_tag.h", gen_tags_enum)
+
+        # Early return - write outputs for CMake dependency tracking
+        if options.output_dependencies:
+            headeronly_fm.write_outputs(
+                "headeronly_generated_headers", options.output_dependencies
+            )
+        return
+
     if "sources" in options.generate:
         gen_source_files(
             native_functions=native_functions,

torchgen/utils.py

@@ -330,8 +330,10 @@ class FileManager:
 def make_file_manager(
     options: Namespace,
     install_dir: str | Path | None = None,
+    template_dir: str | Path | None = None,
 ) -> FileManager:
-    template_dir = os.path.join(options.source_path, "templates")
+    if template_dir is None:
+        template_dir = os.path.join(options.source_path, "templates")
     install_dir = install_dir if install_dir else options.install_dir
     return FileManager(
         install_dir=install_dir,