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

@@ -400,6 +400,7 @@ def get_aten_generated_files(enabled_backends):
         "core/aten_interned_strings.h",
         "core/enum_tag.h",
         "torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.cpp",
+        "torch/headeronly/core/enum_tag.h",
     ] + get_aten_derived_type_srcs(enabled_backends)
 
     # This is tiresome.  A better strategy would be to unconditionally

build.bzl

@@ -93,7 +93,7 @@ def define_targets(rules):
 
     rules.genrule(
         name = "gen_aten",
-        srcs = gen_aten_srcs,
+        srcs = gen_aten_srcs + ["//torch/headeronly:enum_tag_h"],  # Depend on enum_tag.h from torch/headeronly
         outs = gen_aten_outs,
         cmd = gen_aten_cmd,
         tools = ["//torchgen:gen"],

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,47 @@ def define_targets(rules):
 
     rules.genrule(**genrule_args)
 
+    # Generate enum_tag.h from tags.yaml using torchgen
+    # This file is generated in the torch/headeronly subpackage to be at the correct
+    # path (torch/headeronly/core/enum_tag.h) that the ATen forwarding header expects.
+    enum_tag_genrule_args = {
+        "name": "enum_tag_h",
+        "srcs": [
+            "//:aten/src/ATen/native/tags.yaml",
+        ] + rules.glob(["templates/*"]),
+        "outs": ["core/enum_tag.h"],
+        "cmd": " && ".join([
+            # Create a temporary source directory structure
+            "SRCDIR=$$(mktemp -d)",
+            "mkdir -p $$SRCDIR/aten/src/ATen/native",
+            "cp $(location //:aten/src/ATen/native/tags.yaml) $$SRCDIR/aten/src/ATen/native/",
+            # Create template directory
+            "mkdir -p $$SRCDIR/torch/headeronly/templates",
+            "cp $(location templates/enum_tag.h) $$SRCDIR/torch/headeronly/templates/",
+            # Run torchgen to generate only headeronly files
+            "$(execpath //torchgen:gen) " +
+            "--source-path $$SRCDIR/aten/src/ATen " +
+            "--install_dir $$(dirname $$(dirname $$(@D))) " +
+            "--generate headeronly",
+            # Clean up
+            "rm -rf $$SRCDIR",
+        ]),
+        "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",
+            ":enum_tag_h",  # Reference the locally generated 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

@@ -2072,6 +2072,7 @@ def gen_headers(
     cpu_fm: FileManager,
     device_fms: dict[str, FileManager],
     ops_fm: FileManager,
+    headeronly_fm: FileManager,
     dispatch_keys: Sequence[DispatchKey],
     functions_keys: set[DispatchKey],
     rocm: bool,
@@ -2191,10 +2192,15 @@ def gen_headers(
 
     core_fm.write("aten_interned_strings.h", gen_aten_interned_strings)
 
+    # Generate enum_tag.h in torch/headeronly/core
     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)
+    assert headeronly_fm is not None, "headeronly_fm must be provided to gen_headers"
+    headeronly_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(
@@ -2819,7 +2825,7 @@ def main() -> None:
         "--generate",
         type=str,
         nargs="*",
-        choices=["headers", "sources", "declarations_yaml"],
+        choices=["headers", "sources", "declarations_yaml", "headeronly"],
         default=["headers", "sources", "declarations_yaml"],
         help="Generate only a subset of files",
     )
@@ -2945,12 +2951,32 @@ def main() -> None:
     aoti_install_dir = f"{options.aoti_install_dir}"
     Path(aoti_install_dir).mkdir(parents=True, exist_ok=True)
 
+    # Create directory and file manager for torch/headeronly
+    # Derive headeronly_install_dir from install_dir (build output)
+    # Going up 3 levels from build/aten/src/ATen -> build, then to torch/headeronly/core
+    install_dir_parent = os.path.dirname(
+        os.path.dirname(os.path.dirname(options.install_dir))
+    )
+    headeronly_install_dir = f"{install_dir_parent}/torch/headeronly/core"
+    Path(headeronly_install_dir).mkdir(parents=True, exist_ok=True)
+    # Derive headeronly_template_dir from source_path (source tree)
+    # Going up 3 levels from aten/src/ATen -> root, then to torch/headeronly/templates
+    source_path_parent = os.path.dirname(
+        os.path.dirname(os.path.dirname(options.source_path))
+    )
+    headeronly_template_dir = f"{source_path_parent}/torch/headeronly/templates"
+
     core_fm = make_file_manager(options=options, install_dir=core_install_dir)
     cpu_fm = make_file_manager(options=options)
     cpu_vec_fm = make_file_manager(options=options)
     cuda_fm = make_file_manager(options=options)
     ops_fm = make_file_manager(options=options, install_dir=ops_install_dir)
     aoti_fm = make_file_manager(options=options, install_dir=aoti_install_dir)
+    headeronly_fm = make_file_manager(
+        options=options,
+        install_dir=headeronly_install_dir,
+        template_dir=headeronly_template_dir,
+    )
     device_fms = {"cuda": cuda_fm}
     if options.xpu:
         device_fms["xpu"] = make_file_manager(options=options)
@@ -3004,12 +3030,20 @@ def main() -> None:
             cpu_fm=cpu_fm,
             device_fms=device_fms,
             ops_fm=ops_fm,
+            headeronly_fm=headeronly_fm,
             dispatch_keys=dispatch_keys,
             functions_keys=functions_keys,
             rocm=options.rocm,
             per_operator_headers=options.per_operator_headers,
         )
 
+    if "headeronly" in options.generate:
+        # Generate only the headeronly files (enum_tag.h)
+        def gen_tags_enum() -> dict[str, str]:
+            return {"enum_of_valid_tags": (",\n".join(sorted(valid_tags)))}
+
+        headeronly_fm.write("enum_tag.h", gen_tags_enum)
+
     if "declarations_yaml" in options.generate:
         gen_declarations_yaml(native_functions=native_functions, cpu_fm=cpu_fm)
 

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,