From 4276e7da91c1ae06864d6c568403c726d9ee4650 Mon Sep 17 00:00:00 2001
From: Michael Andreas Dagitses <mikeyd@fb.com>
Date: Fri, 17 Jun 2022 05:29:15 -0700
Subject: [PATCH] switch Bazel to common aten codegen

We can also clean up the other rule generation. This is removing
support for per-operator headers right now. We don't take advantage of
that internally, so it's better to have the Bazel build compatible
with our internal builds.

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

**NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D36489359/)!

[ghstack-poisoned]
---
 BUILD.bazel | 21 ++-------------------
 aten.bzl    | 40 ----------------------------------------
 build.bzl   |  3 ---
 3 files changed, 2 insertions(+), 62 deletions(-)

diff --git a/BUILD.bazel b/BUILD.bazel
index ea15f0ec3450de..74d2a1a51421ba 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -6,7 +6,7 @@ load("//third_party:substitution.bzl", "header_template_rule")
 load("//:tools/bazel.bzl", "rules")
 load("//tools/rules:cu.bzl", "cu_library")
 load("//tools/config:defs.bzl", "if_cuda")
-load("//:aten.bzl", "intern_build_aten_ops", "generate_aten")
+load("//:aten.bzl", "intern_build_aten_ops")
 load(":build.bzl", "define_targets", "GENERATED_AUTOGRAD_CPP", "GENERATED_AUTOGRAD_PYTHON")
 load(":build_variables.bzl", "jit_core_sources", "libtorch_core_sources", "libtorch_cuda_sources", "libtorch_distributed_sources", "libtorch_extra_sources", "libtorch_nvfuser_generated_headers", "libtorch_nvfuser_runtime_sources", "libtorch_python_core_sources", "torch_cpp_srcs", "lazy_tensor_ts_sources")
 load(":ufunc_defs.bzl", "aten_ufunc_generated_cpu_sources", "aten_ufunc_generated_cpu_kernel_sources", "aten_ufunc_generated_cuda_sources")
@@ -22,7 +22,6 @@ COMMON_COPTS = [
     "-DTH_HAVE_THREAD",
     "-DUSE_FBGEMM",
     "-DUSE_DISTRIBUTED",
-    "-DAT_PER_OPERATOR_HEADERS",
     "-DATEN_THREADING=NATIVE",
     "-DNO_CUDNN_DESTROY_HANDLE",
 ] + if_cuda([
@@ -30,8 +29,6 @@ COMMON_COPTS = [
     "-DUSE_CUDNN",
 ])
 
-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",
     "aten/src/ATen/RegisterCPU.cpp",
@@ -93,20 +90,6 @@ generated_cuda_cpp = [
     "aten/src/ATen/RegisterSparseCsrCUDA.cpp",
 ]
 
-generate_aten(
-    name = "generated_aten_cpp",
-    srcs = aten_generation_srcs,
-    outs = (
-        generated_cpu_cpp +
-        generated_cuda_cpp +
-        aten_ufunc_generated_cpu_sources("aten/src/ATen/{}") +
-        aten_ufunc_generated_cpu_kernel_sources("aten/src/ATen/{}") +
-        aten_ufunc_generated_cuda_sources("aten/src/ATen/{}") +
-        ["aten/src/ATen/Declarations.yaml"]
-    ),
-    generator = "//torchgen:gen",
-)
-
 filegroup(
     name = "cpp_generated_code",
     data = [":generate-code"],
@@ -291,7 +274,7 @@ cc_library(
     ],
     ) + [
         ":aten_src_ATen_config",
-        ":generated_aten_cpp",
+        ":gen_aten",
     ],
     includes = [
         "aten/src",
diff --git a/aten.bzl b/aten.bzl
index 9dd703717c0748..14dc5a9b247e60 100644
--- a/aten.bzl
+++ b/aten.bzl
@@ -1,4 +1,3 @@
-load("@bazel_skylib//lib:paths.bzl", "paths")
 load("@rules_cc//cc:defs.bzl", "cc_library")
 
 CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX2"]
@@ -55,42 +54,3 @@ def intern_build_aten_ops(copts, deps, extra_impls):
         deps = [":ATen_CPU_" + cpu_capability for cpu_capability in CPU_CAPABILITY_NAMES],
         linkstatic = 1,
     )
-
-def generate_aten_impl(ctx):
-    # Declare the entire ATen/ops/ directory as an output
-    ops_dir = ctx.actions.declare_directory("aten/src/ATen/ops")
-    outputs = [ops_dir] + ctx.outputs.outs
-
-    install_dir = paths.dirname(ops_dir.path)
-    tool_inputs, tool_inputs_manifest = ctx.resolve_tools(tools = [ctx.attr.generator])
-    ctx.actions.run_shell(
-        outputs = outputs,
-        inputs = ctx.files.srcs,
-        command = ctx.executable.generator.path + " $@",
-        arguments = [
-            "--source-path",
-            "aten/src/ATen",
-            "--per-operator-headers",
-            "--install_dir",
-            install_dir,
-        ],
-        tools = tool_inputs,
-        input_manifests = tool_inputs_manifest,
-        use_default_shell_env = True,
-        mnemonic = "GenerateAten",
-    )
-    return [DefaultInfo(files = depset(outputs))]
-
-generate_aten = rule(
-    implementation = generate_aten_impl,
-    attrs = {
-        "generator": attr.label(
-            executable = True,
-            allow_files = True,
-            mandatory = True,
-            cfg = "exec",
-        ),
-        "outs": attr.output_list(),
-        "srcs": attr.label_list(allow_files = True),
-    },
-)
diff --git a/build.bzl b/build.bzl
index a2e84f13c5c0b2..5603f196a74231 100644
--- a/build.bzl
+++ b/build.bzl
@@ -71,9 +71,6 @@ def define_targets(rules):
         tools = ["//torchgen:gen"],
         outs = ["aten/src/ATen/" + out for out in gen_aten_outs],
         cmd = gen_aten_cmd,
-        # This currently clashes with the existing Bazel generated
-        # files.
-        tags = ["-bazel"],
     )
 
     rules.genrule(