xref: /aosp_15_r20/external/pytorch/third_party/tensorpipe.BUILD (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1load("@rules_cc//cc:defs.bzl", "cc_library")
2load("@pytorch//third_party:substitution.bzl", "header_template_rule")
3
4LIBUV_COMMON_SRCS = [
5    "third_party/libuv/src/fs-poll.c",
6    "third_party/libuv/src/idna.c",
7    "third_party/libuv/src/inet.c",
8    "third_party/libuv/src/random.c",
9    "third_party/libuv/src/strscpy.c",
10    "third_party/libuv/src/threadpool.c",
11    "third_party/libuv/src/timer.c",
12    "third_party/libuv/src/uv-common.c",
13    "third_party/libuv/src/uv-data-getter-setters.c",
14    "third_party/libuv/src/version.c",
15]
16
17LIBUV_POSIX_SRCS = [
18    "third_party/libuv/src/unix/async.c",
19    "third_party/libuv/src/unix/core.c",
20    "third_party/libuv/src/unix/dl.c",
21    "third_party/libuv/src/unix/fs.c",
22    "third_party/libuv/src/unix/getaddrinfo.c",
23    "third_party/libuv/src/unix/getnameinfo.c",
24    "third_party/libuv/src/unix/loop.c",
25    "third_party/libuv/src/unix/loop-watcher.c",
26    "third_party/libuv/src/unix/pipe.c",
27    "third_party/libuv/src/unix/poll.c",
28    "third_party/libuv/src/unix/process.c",
29    "third_party/libuv/src/unix/random-devurandom.c",
30    "third_party/libuv/src/unix/signal.c",
31    "third_party/libuv/src/unix/stream.c",
32    "third_party/libuv/src/unix/tcp.c",
33    "third_party/libuv/src/unix/thread.c",
34    "third_party/libuv/src/unix/tty.c",
35    "third_party/libuv/src/unix/udp.c",
36]
37
38LIBUV_LINUX_SRCS = LIBUV_POSIX_SRCS + [
39    "third_party/libuv/src/unix/proctitle.c",
40    "third_party/libuv/src/unix/linux-core.c",
41    "third_party/libuv/src/unix/linux-inotify.c",
42    "third_party/libuv/src/unix/linux-syscalls.c",
43    "third_party/libuv/src/unix/procfs-exepath.c",
44    "third_party/libuv/src/unix/random-getrandom.c",
45    "third_party/libuv/src/unix/random-sysctl-linux.c",
46]
47
48cc_library(
49    name = "libuv",
50    srcs = LIBUV_COMMON_SRCS + LIBUV_LINUX_SRCS,
51    includes = [
52        "third_party/libuv/include",
53        "third_party/libuv/src",
54    ],
55    hdrs = glob(
56        [
57            "third_party/libuv/include/*.h",
58            "third_party/libuv/include/uv/*.h",
59            "third_party/libuv/src/*.h",
60            "third_party/libuv/src/unix/*.h",
61        ],
62    ),
63    visibility = ["//visibility:public"],
64)
65
66cc_library(
67    name = "libnop",
68    srcs = [],
69    includes = ["third_party/libnop/include"],
70    hdrs = glob(["third_party/libnop/include/**/*.h"]),
71)
72
73header_template_rule(
74    name = "tensorpipe_cpu_config_header",
75    src = "tensorpipe/config.h.in",
76    out = "tensorpipe/config.h",
77    substitutions = {
78        "#cmakedefine01 TENSORPIPE_HAS_SHM_TRANSPORT": "#define TENSORPIPE_HAS_SHM_TRANSPORT 1",
79        "#cmakedefine01 TENSORPIPE_HAS_IBV_TRANSPORT": "#define TENSORPIPE_HAS_IBV_TRANSPORT 1",
80        "#cmakedefine01 TENSORPIPE_HAS_CMA_CHANNEL": "#define TENSORPIPE_HAS_CMA_CHANNEL 1",
81    },
82)
83
84header_template_rule(
85    name = "tensorpipe_cuda_config_header",
86    src = "tensorpipe/config_cuda.h.in",
87    out = "tensorpipe/config_cuda.h",
88    substitutions = {
89        "#cmakedefine01 TENSORPIPE_HAS_CUDA_IPC_CHANNEL": "#define TENSORPIPE_HAS_CUDA_IPC_CHANNEL 1",
90        "#cmakedefine01 TENSORPIPE_HAS_CUDA_GDR_CHANNEL": "#define TENSORPIPE_HAS_CUDA_GDR_CHANNEL 1",
91    },
92)
93
94# We explicitly list the CUDA headers & sources, and we consider everything else
95# as CPU (using a catch-all glob). This is both because there's fewer CUDA files
96# (thus making it easier to list them exhaustively) and because it will make it
97# more likely to catch a misclassified file: if we forget to mark a file as CUDA
98# we'll try to build it on CPU and that's likely to fail.
99
100TENSORPIPE_CUDA_HEADERS = [
101    "tensorpipe/tensorpipe_cuda.h",
102    "tensorpipe/channel/cuda_basic/*.h",
103    "tensorpipe/channel/cuda_gdr/*.h",
104    "tensorpipe/channel/cuda_ipc/*.h",
105    "tensorpipe/channel/cuda_xth/*.h",
106    "tensorpipe/common/cuda.h",
107    "tensorpipe/common/cuda_buffer.h",
108    "tensorpipe/common/cuda_lib.h",
109    "tensorpipe/common/cuda_loop.h",
110    "tensorpipe/common/nvml_lib.h",
111]
112
113TENSORPIPE_CUDA_SOURCES = [
114    "tensorpipe/channel/cuda_basic/*.cc",
115    "tensorpipe/channel/cuda_gdr/*.cc",
116    "tensorpipe/channel/cuda_ipc/*.cc",
117    "tensorpipe/channel/cuda_xth/*.cc",
118    "tensorpipe/common/cuda_buffer.cc",
119    "tensorpipe/common/cuda_loop.cc",
120]
121
122TENSORPIPE_CPU_HEADERS = glob(
123    [
124        "tensorpipe/*.h",
125        "tensorpipe/channel/*.h",
126        "tensorpipe/channel/*/*.h",
127        "tensorpipe/common/*.h",
128        "tensorpipe/core/*.h",
129        "tensorpipe/transport/*.h",
130        "tensorpipe/transport/*/*.h",
131    ],
132    exclude=TENSORPIPE_CUDA_HEADERS)
133
134TENSORPIPE_CPU_SOURCES = glob(
135    [
136        "tensorpipe/*.cc",
137        "tensorpipe/channel/*.cc",
138        "tensorpipe/channel/*/*.cc",
139        "tensorpipe/common/*.cc",
140        "tensorpipe/core/*.cc",
141        "tensorpipe/transport/*.cc",
142        "tensorpipe/transport/*/*.cc",
143    ],
144    exclude=TENSORPIPE_CUDA_SOURCES)
145
146cc_library(
147    name = "tensorpipe_cpu",
148    srcs = TENSORPIPE_CPU_SOURCES,
149    hdrs = TENSORPIPE_CPU_HEADERS + [":tensorpipe_cpu_config_header"],
150    includes = [
151        ".",
152    ],
153    copts = [
154        "-std=c++14",
155    ],
156    visibility = ["//visibility:public"],
157    deps = [
158        ":libnop",
159        ":libuv",
160    ],
161)
162
163cc_library(
164    name = "tensorpipe_cuda",
165    srcs = glob(TENSORPIPE_CUDA_SOURCES),
166    hdrs = glob(TENSORPIPE_CUDA_HEADERS) + [":tensorpipe_cuda_config_header"],
167    includes = [
168        ".",
169    ],
170    copts = [
171        "-std=c++14",
172    ],
173    visibility = ["//visibility:public"],
174    deps = [
175        ":tensorpipe_cpu",
176        "@cuda",
177    ],
178)
179