• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Description: Utilities for TPU Operations
2
3load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
4load(
5    "//tensorflow:tensorflow.bzl",
6    "if_windows",
7)
8
9package(
10    default_visibility = [
11        "//tensorflow/compiler/mlir/tensorflow:__subpackages__",
12        "//tensorflow/compiler/tf2xla/kernels:__subpackages__",
13        "//tensorflow/compiler/xrt:__subpackages__",
14        "//tensorflow/core/profiler/internal/tpu:__subpackages__",
15        "//tensorflow/core/tpu:__subpackages__",
16        "//tensorflow/stream_executor/tpu:__subpackages__",
17    ],
18    licenses = ["notice"],  # Apache 2.0
19)
20
21cc_library(
22    name = "libtftpu_header",
23    hdrs = ["libtftpu.h"],
24    visibility = ["//visibility:public"],
25    deps = [],
26)
27
28cc_library(
29    name = "tpu_embedding_optimization_parameters_utils",
30    srcs = ["tpu_embedding_optimization_parameters_utils.cc"],
31    hdrs = ["tpu_embedding_optimization_parameters_utils.h"],
32    visibility = ["//visibility:public"],
33    deps = [
34        "//tensorflow/compiler/xla:xla_data_proto_cc",
35        "//tensorflow/compiler/xla/service:hlo",
36        "//tensorflow/compiler/xla/service:hlo_proto_cc",
37        "//tensorflow/core:framework",
38        "//tensorflow/core:lib",
39        "//tensorflow/core:lib_proto_parsing",
40        "//tensorflow/core:protos_all_cc",
41        "//tensorflow/core/protobuf/tpu:optimization_parameters_proto_cc",
42        "@com_google_absl//absl/base",
43    ],
44)
45
46cc_library(
47    name = "tpu_embedding_output_layout_utils",
48    srcs = ["tpu_embedding_output_layout_utils.cc"],
49    hdrs = ["tpu_embedding_output_layout_utils.h"],
50    visibility = ["//visibility:public"],
51    deps = [
52        "//tensorflow/core:lib_proto_parsing",
53        "//tensorflow/core:protos_all_cc",
54        "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_cc",
55        "//tensorflow/core/protobuf/tpu:tpu_embedding_output_layout_proto_cc",
56    ],
57)
58
59cc_library(
60    name = "tpu_node_device_util",
61    srcs = ["tpu_node_device_util.cc"],
62    hdrs = ["tpu_node_device_util.h"],
63    visibility = ["//visibility:public"],
64    deps = [
65        "//tensorflow/compiler/tf2xla:tf2xla_util",
66        "//tensorflow/core:lib",
67        "//tensorflow/core:protos_all_cc",
68    ],
69)
70
71cc_library(
72    name = "tpu_compile_interface",
73    srcs = ["tpu_compile_interface.cc"],
74    hdrs = ["tpu_compile_interface.h"],
75    deps = [
76        "//tensorflow/core:lib",
77        "@com_google_absl//absl/strings",
78    ],
79)
80
81cc_library(
82    name = "tpu_defs",
83    srcs = ["tpu_defs.cc"],
84    hdrs = ["tpu_defs.h"],
85    visibility = ["//visibility:public"],
86    deps = ["//tensorflow/core:protos_all_cc"],
87)
88
89cc_library(
90    name = "tpu_configuration",
91    srcs = ["tpu_configuration.cc"],
92    hdrs = ["tpu_configuration.h"],
93    deps = ["//tensorflow/core:framework"],
94)
95
96cc_library(
97    name = "tpu_init_mode",
98    srcs = ["tpu_init_mode.cc"],
99    hdrs = ["tpu_init_mode.h"],
100    deps = [
101        "//tensorflow/core:lib",
102    ],
103)
104
105cc_library(
106    name = "tpu_initializer_helper",
107    srcs = ["tpu_initializer_helper.cc"],
108    hdrs = ["tpu_initializer_helper.h"],
109    deps = [
110        "//tensorflow/core/platform:logging",
111        "@com_google_absl//absl/strings",
112        "@com_google_absl//absl/synchronization",
113    ],
114)
115
116cc_library(
117    name = "tpu_api",
118    srcs = ["tpu_api.cc"],
119    hdrs = ["tpu_api.h"],
120    visibility = ["//visibility:public"],
121    deps = [
122        ":libtftpu_header",
123        ":tpu_executor_api",
124        ":tpu_ops_c_api_hdrs",
125        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
126    ],
127)
128
129cc_library(
130    name = "tpu_executor_api",
131    srcs = ["tpu_executor_api.cc"],
132    hdrs = ["tpu_executor_api.h"],
133    deps = [
134        ":libtftpu_header",
135        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
136    ],
137)
138
139cc_library(
140    name = "tpu_api_dlsym_initializer",
141    srcs = if_windows(
142        ["tpu_api_dlsym_initializer_windows.cc"],
143        otherwise = ["tpu_api_dlsym_initializer.cc"],
144    ),
145    hdrs = ["tpu_api_dlsym_initializer.h"],
146    visibility = ["//visibility:public"],
147    deps = [
148        ":libtftpu_header",
149        ":tpu_api",
150        ":tpu_api_dlsym_set_fn",
151        ":tpu_executor_init_fns",
152        ":tpu_initializer_helper",
153        ":tpu_library_init_fns",
154        ":tpu_ops_c_api_hdrs",
155        "//tensorflow/core:lib",
156        "//tensorflow/stream_executor/tpu:tpu_executor",
157        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
158    ],
159    # Always link this in, because even if we don't use it directly we want its
160    #static initializers to dynamically load API symbols exported from libtpu.so
161    alwayslink = True,
162)
163
164# This is an alternative to "tpu_api_dlsym_initializer" that only initializes
165# methods needed for the base TPU executor APIs (and thus has fewer deps). Do
166# not link in both this and "tpu_api_dlsym_initializer".
167cc_library(
168    name = "tpu_executor_dlsym_initializer",
169    srcs = if_windows(
170        ["tpu_executor_dlsym_initializer_windows.cc"],
171        otherwise = ["tpu_executor_dlsym_initializer.cc"],
172    ),
173    visibility = ["//visibility:public"],
174    deps = [
175        ":tpu_api_dlsym_set_fn",
176        ":tpu_executor_init_fns",
177        ":tpu_initializer_helper",
178        "//tensorflow/compiler/xla/service:tpu_computation_placer",
179        "//tensorflow/core:lib",
180        "//tensorflow/stream_executor/tpu:tpu_executor",
181        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
182    ],
183    alwayslink = True,
184)
185
186cc_library(
187    name = "tpu_api_dlsym_set_fn",
188    hdrs = ["tpu_api_dlsym_set_fn.h"],
189    visibility = ["//visibility:public"],
190)
191
192cc_library(
193    name = "tpu_library_init_fns",
194    hdrs = ["tpu_library_init_fns.inc"],
195    visibility = ["//visibility:public"],
196    deps = [":tpu_executor_init_fns"],
197)
198
199cc_library(
200    name = "tpu_executor_init_fns",
201    hdrs = ["tpu_executor_init_fns.inc"],
202    visibility = ["//visibility:public"],
203)
204
205cc_library(
206    name = "virtual_device",
207    srcs = ["virtual_device.cc"],
208    hdrs = ["virtual_device.h"],
209    visibility = ["//visibility:public"],
210    deps = [
211        "//tensorflow/core:core_cpu",
212        "//tensorflow/core:protos_all_cc",
213    ],
214)
215
216cc_library(
217    name = "tpu_execute",
218    srcs = ["tpu_execute.cc"],
219    hdrs = ["tpu_execute.h"],
220    deps = [
221        ":tpu_api",
222        "//tensorflow/compiler/xla:executable_run_options",
223        "//tensorflow/compiler/xla:shape_layout",
224        "//tensorflow/compiler/xla:shape_util",
225        "//tensorflow/compiler/xla:status",
226        "//tensorflow/compiler/xla:status_macros",
227        "//tensorflow/compiler/xla:statusor",
228        "//tensorflow/compiler/xla:util",
229        "//tensorflow/compiler/xla:xla_data_proto_cc",
230        "//tensorflow/compiler/xla/service:computation_layout",
231        "//tensorflow/compiler/xla/service:computation_placer",
232        "//tensorflow/compiler/xla/service:executable",
233        "//tensorflow/compiler/xla/service:hlo",
234        "//tensorflow/compiler/xla/service:hlo_module_config",
235        "//tensorflow/compiler/xla/service:hlo_proto_cc",
236        "//tensorflow/compiler/xla/service:maybe_owning_device_memory",
237        "//tensorflow/compiler/xla/service:transfer_manager",
238        "//tensorflow/core:framework",
239        "//tensorflow/core:lib",
240        "//tensorflow/core/profiler/lib:traceme",
241        "//tensorflow/core/tpu:tpu_ops_c_api_hdrs",
242        "//tensorflow/core/tpu/kernels:tpu_executable_info_proto_cc",
243        "//tensorflow/stream_executor:device_memory",
244        "//tensorflow/stream_executor:stream",
245        "//tensorflow/stream_executor/lib",
246        "//tensorflow/stream_executor/tpu:c_api_conversions",
247        "//tensorflow/stream_executor/tpu:status_helper",
248        "//tensorflow/stream_executor/tpu:tpu_executable",
249        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
250        "//tensorflow/stream_executor/tpu:tpu_node_context",
251        "//tensorflow/stream_executor/tpu:tpu_platform_interface",
252        "@com_google_absl//absl/base",
253        "@com_google_absl//absl/memory",
254    ],
255)
256
257cc_library(
258    name = "tpu_on_demand_compiler",
259    srcs = ["tpu_on_demand_compiler.cc"],
260    visibility = ["//visibility:public"],
261    deps = [
262        "//tensorflow/compiler/xla:shape_util",
263        "//tensorflow/compiler/xla:util",
264        "//tensorflow/compiler/xla:xla_data_proto_cc",
265        "//tensorflow/compiler/xla/service:compiler",
266        "//tensorflow/compiler/xla/service:executable",
267        "//tensorflow/compiler/xla/service:hlo",
268        "//tensorflow/compiler/xla/service:hlo_cost_analysis",
269        "//tensorflow/compiler/xla/service:hlo_module_group",
270        "//tensorflow/compiler/xla/service:shaped_buffer",
271        "//tensorflow/stream_executor:device_memory_allocator",
272        "//tensorflow/stream_executor/tpu:c_api_conversions",
273        "//tensorflow/stream_executor/tpu:c_api_decl",
274        "//tensorflow/stream_executor/tpu:proto_helper",
275        "//tensorflow/stream_executor/tpu:status_helper",
276        "//tensorflow/stream_executor/tpu:tpu_executable_interface",
277        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
278        "//tensorflow/stream_executor/tpu:tpu_executor_hdrs",
279        "//tensorflow/stream_executor/tpu:tpu_platform_id",
280        "@com_google_absl//absl/types:span",
281    ],
282    alwayslink = True,
283)
284
285cc_library(
286    name = "tpu_runtime",
287    srcs = [],
288    visibility = ["//visibility:public"],
289    deps = [
290        ":tpu_api_dlsym_initializer",
291        "//tensorflow/core/tpu:tpu_on_demand_compiler",
292        "//tensorflow/core/tpu/ops",
293    ],
294)
295
296cc_library(
297    name = "tpu_ops_c_api_hdrs",
298    srcs = [],
299    hdrs = ["tpu_ops_c_api.h"],
300    visibility = ["//visibility:public"],
301    deps = [
302        ":libtftpu_header",
303        "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
304        "//tensorflow/stream_executor/tpu:c_api_decl",
305        "//tensorflow/stream_executor/tpu:proto_helper",
306    ],
307    alwayslink = True,
308)
309
310cc_library(
311    name = "tpu_fingerprint_utils",
312    srcs = ["tpu_fingerprint_utils.cc"],
313    hdrs = ["tpu_fingerprint_utils.h"],
314    deps = [
315        ":tpu_compile_interface",
316        "//tensorflow/compiler/xla:status_macros",
317        "//tensorflow/core:framework",
318        "//tensorflow/core:lib_internal_impl",
319        "//tensorflow/core/lib/core:status",
320    ],
321)
322
323cc_library(
324    name = "tpu_model_server_initializer",
325    srcs = ["tpu_model_server_initializer.cc"],
326    hdrs = ["tpu_model_server_initializer.h"],
327    visibility = ["//visibility:public"],
328    deps = [
329        ":libtftpu_header",
330        ":tpu_api",
331        ":tpu_api_dlsym_set_fn",
332        ":tpu_executor_init_fns",
333        ":tpu_initializer_helper",
334        ":tpu_library_init_fns",
335        ":tpu_ops_c_api_hdrs",
336        "//tensorflow/core:lib",
337        "//tensorflow/stream_executor/tpu:tpu_executor",
338        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
339    ],
340    alwayslink = True,
341)
342