Revert "Revert "Allow specifying tags for aten operators in native_fu…

…nctions.yaml""

This reverts commit ea44645.

Pull Request resolved: pytorch#76456

Approved by: https://github.com/osalpekar

.circleci/scripts/cpp_doc_push_script.sh

@@ -65,7 +65,8 @@ cp torch/_utils_internal.py tools/shared
 
 # Generate PyTorch files
 time python tools/setup_helpers/generate_code.py \
-  --native-functions-path aten/src/ATen/native/native_functions.yaml
+  --native-functions-path aten/src/ATen/native/native_functions.yaml \
+  --tags-path aten/src/ATen/native/tags.yaml
 
 # Build the docs
 pushd docs/cpp

.github/workflows/lint.yml

@@ -33,6 +33,7 @@ jobs:
           python3 -m tools.generate_torch_version --is_debug=false
           python3 -m tools.pyi.gen_pyi \
             --native-functions-path aten/src/ATen/native/native_functions.yaml \
+            --tags-path aten/src/ATen/native/tags.yaml \
             --deprecated-functions-path "tools/autograd/deprecated.yaml"
 
       - name: Run lintrunner on all files

.jenkins/pytorch/codegen-test.sh

@@ -38,19 +38,22 @@ mkdir -p "$OUT"/pyi/torch/_C
 mkdir -p "$OUT"/pyi/torch/nn
 python -m tools.pyi.gen_pyi \
   --native-functions-path aten/src/ATen/native/native_functions.yaml \
+  --tags-path aten/src/ATen/native/tags.yaml \
   --deprecated-functions-path tools/autograd/deprecated.yaml \
   --out "$OUT"/pyi
 
 # autograd codegen (called by torch codegen but can run independently)
 python -m tools.autograd.gen_autograd \
   "$OUT"/torch/share/ATen/Declarations.yaml \
   aten/src/ATen/native/native_functions.yaml \
+  aten/src/ATen/native/tags.yaml \
   "$OUT"/autograd \
   tools/autograd
 
 # annotated_fn_args codegen (called by torch codegen but can run independently)
 mkdir -p "$OUT"/annotated_fn_args
 python -m tools.autograd.gen_annotated_fn_args \
   aten/src/ATen/native/native_functions.yaml \
+  aten/src/ATen/native/tags.yaml \
   "$OUT"/annotated_fn_args \
   tools/autograd

BUILD.bazel

@@ -29,7 +29,7 @@ COMMON_COPTS = [
     "-DUSE_CUDNN",
 ])
 
-aten_generation_srcs = ["aten/src/ATen/native/native_functions.yaml"] + glob(["aten/src/ATen/templates/**"])
+aten_generation_srcs = ["aten/src/ATen/native/native_functions.yaml"] + ["aten/src/ATen/native/tags.yaml"] + glob(["aten/src/ATen/templates/**"])
 
 generated_cpu_cpp = [
     "aten/src/ATen/RegisterBackendSelect.cpp",
@@ -152,6 +152,7 @@ genrule(
     name = "all_generated_code",
     srcs = [
         "aten/src/ATen/native/native_functions.yaml",
+        "aten/src/ATen/native/tags.yaml",
         "aten/src/ATen/native/ts_native_functions.yaml",
         "torch/csrc/lazy/core/shape_inference.h",
         "torch/csrc/lazy/ts_backend/ts_native_functions.cpp",
@@ -161,7 +162,7 @@ genrule(
         "aten/src/ATen/templates/LazyIr.h",
     ],
     outs = libtorch_cpp_generated_sources + libtorch_python_generated_sources,
-    cmd = "$(location //tools/setup_helpers:generate_code) --install_dir `dirname $(location torch/csrc/autograd/generated/variable_factories.h)`/../.. --native-functions-path $(location aten/src/ATen/native/native_functions.yaml) --gen_lazy_ts_backend",
+    cmd = "$(location //tools/setup_helpers:generate_code) --install_dir `dirname $(location torch/csrc/autograd/generated/variable_factories.h)`/../.. --native-functions-path $(location aten/src/ATen/native/native_functions.yaml) --tags-path $(location aten/src/ATen/native/tags.yaml) --gen_lazy_ts_backend",
     tools = ["//tools/setup_helpers:generate_code"],
 )
 

CODEOWNERS

@@ -14,6 +14,7 @@
 /test/allowlist_for_publicAPI.json @albanD @anjali411
 /docs/source/conf.py @albanD
 /aten/src/ATen/native/native_functions.yaml @bdhirsh
+/aten/src/ATen/native/tags.yaml @anjali411
 
 # Tensorpipe RPC Agent.
 /torch/csrc/distributed/rpc/tensorpipe_agent.cpp @jiayisuse @osalpekar @lw @beauby

caffe2/CMakeLists.txt

@@ -417,13 +417,15 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     COMMAND
     "${PYTHON_EXECUTABLE}" tools/setup_helpers/generate_code.py
       --native-functions-path "aten/src/ATen/native/native_functions.yaml"
+      --tags-path "aten/src/ATen/native/tags.yaml"
       $<$<BOOL:${INTERN_DISABLE_AUTOGRAD}>:--disable-autograd>
       $<$<BOOL:${SELECTED_OP_LIST}>:--selected-op-list-path="${SELECTED_OP_LIST}">
       --force_schema_registration
       --gen_lazy_ts_backend
       ${GEN_PER_OPERATOR_FLAG}
     DEPENDS
     "${TORCH_ROOT}/aten/src/ATen/native/native_functions.yaml"
+    "${TORCH_ROOT}/aten/src/ATen/native/tags.yaml"
     "${TORCH_ROOT}/aten/src/ATen/native/ts_native_functions.yaml"
     "${TORCH_ROOT}/torch/csrc/lazy/core/shape_inference.h"
     "${TORCH_ROOT}/torch/csrc/lazy/ts_backend/ts_native_functions.cpp"

cmake/Codegen.cmake

@@ -136,6 +136,7 @@ if(INTERN_BUILD_ATEN_OPS)
         COMMAND ${GEN_UNBOXING_COMMAND_sources}
         DEPENDS ${all_unboxing_script} ${sources_templates}
         ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/native_functions.yaml
+        ${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/tags.yaml
         WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
     )
   else() # Otherwise do not generate or include sources into build.
@@ -210,6 +211,7 @@ if(INTERN_BUILD_ATEN_OPS)
       COMMAND ${GEN_COMMAND_${gen_type}}
       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
       WORKING_DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/..
     )
   endforeach()

docs/cpp/source/check-doxygen.sh

@@ -19,8 +19,8 @@ cp torch/_utils_internal.py tools/shared
 python -m torchgen.gen
 
 python tools/setup_helpers/generate_code.py                 \
-  --native-functions-path aten/src/ATen/native/native_functions.yaml
-
+  --native-functions-path aten/src/ATen/native/native_functions.yaml \
+  --tags-path aten/src/ATen/native/tags.yaml
 popd
 
 # Run doxygen and log all output.

tools/autograd/gen_annotated_fn_args.py

@@ -5,6 +5,7 @@
 
 python -m tools.autograd.gen_annotated_fn_args \
        aten/src/ATen/native/native_functions.yaml \
+       aten/src/ATen/native/tags.yaml \
        $OUTPUT_DIR \
        tools/autograd
 
@@ -36,8 +37,12 @@
 )
 
 
-def gen_annotated(native_yaml_path: str, out: str, autograd_dir: str) -> None:
-    native_functions = parse_native_yaml(native_yaml_path).native_functions
+def gen_annotated(
+    native_yaml_path: str, tags_yaml_path: str, out: str, autograd_dir: str
+) -> None:
+    native_functions = parse_native_yaml(
+        native_yaml_path, tags_yaml_path
+    ).native_functions
     mappings = (
         (is_py_torch_function, "torch._C._VariableFunctions"),
         (is_py_nn_function, "torch._C._nn"),
@@ -90,12 +95,13 @@ def main() -> None:
     parser.add_argument(
         "native_functions", metavar="NATIVE", help="path to native_functions.yaml"
     )
+    parser.add_argument("tags", metavar="TAGS", help="path to tags.yaml")
     parser.add_argument("out", metavar="OUT", help="path to output directory")
     parser.add_argument(
         "autograd", metavar="AUTOGRAD", help="path to template directory"
     )
     args = parser.parse_args()
-    gen_annotated(args.native_functions, args.out, args.autograd)
+    gen_annotated(args.native_functions, args.tags, args.out, args.autograd)
 
 
 if __name__ == "__main__":

tools/autograd/gen_autograd.py

@@ -5,6 +5,7 @@
 python -m tools.autograd.gen_autograd \
        build/aten/src/ATen/Declarations.yaml \
        aten/src/ATen/native/native_functions.yaml \
+       aten/src/ATen/native/tags.yaml \
        $OUTPUT_DIR \
        tools/autograd
 
@@ -46,19 +47,20 @@
 
 def gen_autograd(
     native_functions_path: str,
+    tags_path: str,
     out: str,
     autograd_dir: str,
     operator_selector: SelectiveBuilder,
     disable_autograd: bool = False,
 ) -> None:
     # Parse and load derivatives.yaml
     differentiability_infos = load_derivatives(
-        os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path
+        os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path, tags_path
     )
 
     template_path = os.path.join(autograd_dir, "templates")
 
-    native_funcs = parse_native_yaml(native_functions_path).native_functions
+    native_funcs = parse_native_yaml(native_functions_path, tags_path).native_functions
     fns = list(
         sorted(
             filter(
@@ -74,11 +76,11 @@ def gen_autograd(
     # Generate VariableType.h/cpp
     if not disable_autograd:
         gen_variable_type(
-            out, native_functions_path, fns_with_diff_infos, template_path
+            out, native_functions_path, tags_path, fns_with_diff_infos, template_path
         )
 
         gen_inplace_or_view_type(
-            out, native_functions_path, fns_with_diff_infos, template_path
+            out, native_functions_path, tags_path, fns_with_diff_infos, template_path
         )
 
         # operator filter not applied as tracing sources are excluded in selective build
@@ -87,16 +89,17 @@ def gen_autograd(
     gen_autograd_functions_lib(out, differentiability_infos, template_path)
 
     # Generate variable_factories.h
-    gen_variable_factories(out, native_functions_path, template_path)
+    gen_variable_factories(out, native_functions_path, tags_path, template_path)
 
 
 def gen_autograd_python(
     native_functions_path: str,
+    tags_path: str,
     out: str,
     autograd_dir: str,
 ) -> None:
     differentiability_infos = load_derivatives(
-        os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path
+        os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path, tags_path
     )
 
     template_path = os.path.join(autograd_dir, "templates")
@@ -106,21 +109,25 @@ def gen_autograd_python(
 
     # Generate Python bindings
     deprecated_path = os.path.join(autograd_dir, "deprecated.yaml")
-    gen_python_functions.gen(out, native_functions_path, deprecated_path, template_path)
+    gen_python_functions.gen(
+        out, native_functions_path, tags_path, deprecated_path, template_path
+    )
 
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Generate autograd C++ files script")
     parser.add_argument(
         "native_functions", metavar="NATIVE", help="path to native_functions.yaml"
     )
+    parser.add_argument("tags", metavar="NATIVE", help="path to tags.yaml")
     parser.add_argument("out", metavar="OUT", help="path to output directory")
     parser.add_argument(
         "autograd", metavar="AUTOGRAD", help="path to autograd directory"
     )
     args = parser.parse_args()
     gen_autograd(
         args.native_functions,
+        args.tags,
         args.out,
         args.autograd,
         SelectiveBuilder.get_nop_selector(),

tools/autograd/gen_inplace_or_view_type.py

@@ -563,6 +563,7 @@ def gen_inplace_or_view_type_env(
 def gen_inplace_or_view_type(
     out: str,
     native_yaml_path: str,
+    tags_yaml_path: str,
     fns_with_infos: List[NativeFunctionWithDifferentiabilityInfo],
     template_path: str,
 ) -> None:

tools/autograd/gen_python_functions.py

@@ -233,10 +233,16 @@ def is_py_special_function(f: NativeFunction) -> bool:
 
 
 def gen(
-    out: str, native_yaml_path: str, deprecated_yaml_path: str, template_path: str
+    out: str,
+    native_yaml_path: str,
+    tags_yaml_path: str,
+    deprecated_yaml_path: str,
+    template_path: str,
 ) -> None:
     fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
-    native_functions = parse_native_yaml(native_yaml_path).native_functions
+    native_functions = parse_native_yaml(
+        native_yaml_path, tags_yaml_path
+    ).native_functions
     native_functions = list(filter(should_generate_py_binding, native_functions))
 
     methods = load_signatures(native_functions, deprecated_yaml_path, method=True)

tools/autograd/gen_variable_factories.py

@@ -34,8 +34,12 @@ def maybe_optional_type(type: str, is_opt: bool) -> str:
     return maybe_optional_type(qualified_type, is_opt)
 
 
-def gen_variable_factories(out: str, native_yaml_path: str, template_path: str) -> None:
-    native_functions = parse_native_yaml(native_yaml_path).native_functions
+def gen_variable_factories(
+    out: str, native_yaml_path: str, tags_yaml_path: str, template_path: str
+) -> None:
+    native_functions = parse_native_yaml(
+        native_yaml_path, tags_yaml_path
+    ).native_functions
     factory_functions = [fn for fn in native_functions if is_factory_function(fn)]
     fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
     fm.write_with_template(

tools/autograd/gen_variable_type.py

@@ -672,6 +672,7 @@
 def gen_variable_type(
     out: str,
     native_yaml_path: str,
+    tags_yaml_path: str,
     fns_with_diff_infos: List[NativeFunctionWithDifferentiabilityInfo],
     template_path: str,
 ) -> None:

tools/autograd/load_derivatives.py

@@ -72,7 +72,7 @@ def add_view_copy_derivatives(
 
 
 def load_derivatives(
-    derivatives_yaml_path: str, native_yaml_path: str
+    derivatives_yaml_path: str, native_yaml_path: str, tags_yaml_path: str
 ) -> Sequence[DifferentiabilityInfo]:
     # Do some caching as this is a deterministic function
     global _GLOBAL_LOAD_DERIVATIVE_CACHE
@@ -82,7 +82,7 @@ def load_derivatives(
         with open(derivatives_yaml_path, "r") as f:
             definitions = yaml.load(f, Loader=YamlLoader)
 
-        funcs = parse_native_yaml(native_yaml_path).native_functions
+        funcs = parse_native_yaml(native_yaml_path, tags_yaml_path).native_functions
         # From the parsed native functions, separate out the (generated) view_copy functions,
         # so we can generate derivatives for them separately.
         native_functions_with_view_groups = get_grouped_by_view_native_functions(funcs)

tools/jit/gen_unboxing.py

@@ -217,7 +217,8 @@ def main() -> None:
         selector = SelectiveBuilder.get_nop_selector()
 
     native_yaml_path = os.path.join(options.source_path, "native/native_functions.yaml")
-    parsed_yaml = parse_native_yaml(native_yaml_path)
+    tags_yaml_path = os.path.join(options.source_path, "native/tags.yaml")
+    parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
     native_functions, backend_indices = (
         parsed_yaml.native_functions,
         parsed_yaml.backend_indices,

tools/linter/clang_tidy/generate_build_files.py

@@ -58,6 +58,8 @@ def run_autogen() -> None:
             "tools/setup_helpers/generate_code.py",
             "--native-functions-path",
             "aten/src/ATen/native/native_functions.yaml",
+            "--tags-path",
+            "aten/src/ATen/native/tags.yaml",
             "--gen_lazy_ts_backend",
         ]
     )

tools/pyi/gen_pyi.py

@@ -321,7 +321,12 @@ def gen_nn_functional(fm: FileManager) -> None:
     )
 
 
-def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) -> None:
+def gen_pyi(
+    native_yaml_path: str,
+    tags_yaml_path: str,
+    deprecated_yaml_path: str,
+    fm: FileManager,
+) -> None:
     """gen_pyi()
 
     This function generates a pyi file for torch.
@@ -512,7 +517,9 @@ def gen_pyi(native_yaml_path: str, deprecated_yaml_path: str, fm: FileManager) -
             )
         )
 
-    native_functions = parse_native_yaml(native_yaml_path).native_functions
+    native_functions = parse_native_yaml(
+        native_yaml_path, tags_yaml_path
+    ).native_functions
     native_functions = list(filter(should_generate_py_binding, native_functions))
 
     function_signatures = load_signatures(
@@ -900,6 +907,12 @@ def main() -> None:
         default="aten/src/ATen/native/native_functions.yaml",
         help="path to native_functions.yaml",
     )
+    parser.add_argument(
+        "--tags-path",
+        metavar="TAGS",
+        default="aten/src/ATen/native/tags.yaml",
+        help="path to tags.yaml",
+    )
     parser.add_argument(
         "--deprecated-functions-path",
         metavar="DEPRECATED",
@@ -911,7 +924,9 @@ def main() -> None:
     )
     args = parser.parse_args()
     fm = FileManager(install_dir=args.out, template_dir=".", dry_run=False)
-    gen_pyi(args.native_functions_path, args.deprecated_functions_path, fm)
+    gen_pyi(
+        args.native_functions_path, args.tags_path, args.deprecated_functions_path, fm
+    )
 
 
 if __name__ == "__main__":