create //c10/cuda library (#70863)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/70863

ghstack-source-id: 148159368

Test Plan: Ought to be a no-op: rely on CI to validate.

Reviewed By: malfet

Differential Revision: D33367290

fbshipit-source-id: cb550538b9eafaa0117f94077ebd4cb920688881
(cherry picked from commit 077d9578bc)

c10/BUILD.bazel

@@ -1,6 +1,5 @@
 load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
-load("//third_party:substitution.bzl", "header_template_rule")
 load("//tools/config:defs.bzl", "if_cuda")
 
 # The bool_flag targets allow configuring the build from the
@@ -27,28 +26,17 @@ config_setting(
     flag_values = {":use_glog": "true"},
 )
 
-header_template_rule(
-    name = "cuda_cmake_macros_h",
-    src = "cuda/impl/cuda_cmake_macros.h.in",
-    out = "cuda/impl/cuda_cmake_macros.h",
-    substitutions = {
-        "cmakedefine": "define",
-    },
-)
-
 cc_library(
     name = "headers",
     hdrs = [
         "//c10/core:headers",
+        "//c10/cuda:headers",
         "//c10/mobile:headers",
         "//c10/util:headers",
-    ] + glob([
+    ],
-        "cuda/*.h",
-        "cuda/impl/*.h",
-    ]),
     deps = [
-        ":cuda_cmake_macros_h",
         "//c10/core:alignment",
+        "//c10/cuda:Macros",
         "//c10/macros",
     ] + select({
         ":using_gflags": ["@com_github_gflags_gflags//:gflags"],
@@ -62,14 +50,6 @@ cc_library(
 
 cc_library(
     name = "c10",
-    srcs = if_cuda(
-        glob([
-            "cuda/*.cpp",
-            "cuda/impl/*.cpp",
-        ]),
-        [],
-    ),
-    local_defines = ["C10_BUILD_MAIN_LIB"],
     deps = [
         ":headers",
         "//c10/core:CPUAllocator",
@@ -82,7 +62,7 @@ cc_library(
         "//c10/util:base",
         "//c10/util:typeid",
     ] + if_cuda(
-        ["@cuda"],
+        ["//c10/cuda"],
         [],
     ),
     alwayslink = True,

c10/cuda/BUILD.bazel

@@ -0,0 +1,4 @@
+load("//:tools/bazel.bzl", "rules")
+load(":build.bzl", "define_targets")
+
+define_targets(rules = rules)

c10/cuda/build.bzl

@@ -0,0 +1,65 @@
+def define_targets(rules):
+    rules.cc_library(
+        name = "cuda",
+        srcs = rules.glob(
+            [
+                "*.cpp",
+                "impl/*.cpp",
+            ],
+            exclude = [
+                "test/**/*.cpp",
+            ],
+        ),
+        hdrs = rules.glob(
+            [
+                "*.h",
+                "impl/*.h",
+            ],
+            exclude = [
+                "CUDAMacros.h",
+            ],
+        ),
+        # This library uses registration. Don't let registered
+        # entities be removed.
+        alwayslink = True,
+        linkstatic = True,
+        local_defines = ["C10_BUILD_MAIN_LIB"],
+        visibility = ["//visibility:public"],
+        deps = [
+            ":Macros",
+            "@cuda",
+            "//c10/core:base",
+            "//c10/macros",
+            "//c10/util:base",
+        ],
+        target_compatible_with = rules.requires_cuda_enabled(),
+    )
+
+    rules.cc_library(
+        name = "Macros",
+        srcs = [":cuda_cmake_macros"],
+        hdrs = ["CUDAMacros.h"],
+        linkstatic = True,
+        local_defines = ["C10_BUILD_MAIN_LIB"],
+        visibility = ["//visibility:public"],
+    )
+
+    rules.cmake_configure_file(
+        name = "cuda_cmake_macros",
+        src = "impl/cuda_cmake_macros.h.in",
+        out = "impl/cuda_cmake_macros.h",
+        definitions = [],
+    )
+
+    rules.filegroup(
+        name = "headers",
+        srcs = rules.glob(
+            [
+                "*.h",
+                "impl/*.h",
+            ],
+            exclude = [
+            ],
+        ),
+        visibility = ["//c10:__pkg__"],
+    )

tools/bazel.bzl

@@ -1,4 +1,5 @@
 load("@rules_cc//cc:defs.bzl", "cc_library")
+load("@rules_cuda//cuda:defs.bzl", "requires_cuda_enabled")
 load("//c10/macros:cmake_configure_file.bzl", "cmake_configure_file")
 
 # Rules implementation for the Bazel build system. Since the common
@@ -9,5 +10,6 @@ rules = struct(
     cmake_configure_file = cmake_configure_file,
     filegroup = native.filegroup,
     glob = native.glob,
+    requires_cuda_enabled = requires_cuda_enabled,
     select = select,
 )