use ufunc_defs.bzl in BUILD.bazel (#77747)

mikey dagitses · facebook-github-bot · commit fe4a0dc18e6c · 2022-06-16T00:50:55.000-07:00
Summary: Pull Request resolved: #77747 This allows us to remove some duplication in aten.bzl. ghstack-source-id: 159044901 Test Plan: Rely on CI. Reviewed By: seemethere Differential Revision: D36481025 fbshipit-source-id: e17aabd2e246382219b21a1408a03d577ac1fb4f
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -6,9 +6,10 @@ 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", "aten_ufunc_generated_cpu_sources", "aten_ufunc_generated_cpu_kernel_sources", "aten_ufunc_generated_cuda_sources")
+load("//:aten.bzl", "intern_build_aten_ops", "generate_aten")
 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")
 
 define_targets(rules = rules)
 
diff --git a/aten.bzl b/aten.bzl
@@ -1,6 +1,5 @@
 load("@bazel_skylib//lib:paths.bzl", "paths")
 load("@rules_cc//cc:defs.bzl", "cc_library")
-load(":build_variables.bzl", "aten_ufunc_headers")
 
 CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX2"]
 CAPABILITY_COMPILER_FLAGS = {
@@ -95,32 +94,3 @@ generate_aten = rule(
         "srcs": attr.label_list(allow_files = True),
     },
 )
-
-# copy pasted from ufunc_defs.bzl, as ufuncs_defs.bzl cannot be included
-# from BUILD.bazel because it has a directory relative load, and Bazel
-# always load from workspace root.  The "correct" fix would be to move
-# build_variables.bzl to the top level but I don't have time to do this at
-# the moment.
-
-aten_ufunc_names = [
-    paths.split_extension(paths.basename(h))[0]
-    for h in aten_ufunc_headers
-]
-
-def aten_ufunc_generated_cpu_sources(gencode_pattern = "{}"):
-    return [gencode_pattern.format(name) for name in [
-        "UfuncCPU_{}.cpp".format(n)
-        for n in aten_ufunc_names
-    ]]
-
-def aten_ufunc_generated_cpu_kernel_sources(gencode_pattern = "{}"):
-    return [gencode_pattern.format(name) for name in [
-        "UfuncCPUKernel_{}.cpp".format(n)
-        for n in aten_ufunc_names
-    ]]
-
-def aten_ufunc_generated_cuda_sources(gencode_pattern = "{}"):
-    return [gencode_pattern.format(name) for name in [
-        "UfuncCUDA_{}.cu".format(n)
-        for n in aten_ufunc_names
-    ]]
