• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
2load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
3load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependency")
4load(
5    "//tensorflow/compiler/xla/python/tpu_driver:platform/external/tools.bzl",
6    "external_deps",
7    "go_grpc_library",
8)
9
10licenses(["notice"])  # Apache 2.0
11
12package(default_visibility = ["//visibility:public"])
13
14tf_proto_library(
15    name = "tpu_driver_proto",
16    srcs = ["tpu_driver.proto"],
17    cc_api_version = 2,
18    protodeps = [],
19)
20
21tf_proto_library(
22    name = "tpu_service_proto",
23    srcs = ["tpu_service.proto"],
24    has_services = 1,
25    cc_api_version = 2,
26    cc_grpc_version = 1,
27    protodeps = [
28        ":tpu_driver_proto",
29        "//tensorflow/compiler/xla:xla_data_proto",
30        "//tensorflow/compiler/xla/service:hlo_proto",
31    ],
32    use_grpc_namespace = True,
33)
34
35cc_library(
36    name = "tpu_driver",
37    srcs = [
38        "tpu_driver.cc",
39    ],
40    hdrs = [
41        "event_id.h",
42        "platform/external/compat.h",
43        "tpu_driver.h",
44    ],
45    deps = [
46        "@com_google_absl//absl/types:optional",
47        "//tensorflow/core/platform:logging",
48        "//tensorflow/compiler/xla:status",
49        "//tensorflow/compiler/xla:statusor",
50        "//tensorflow/compiler/xla:util",
51        "//tensorflow/compiler/xla:xla_data_proto_cc",
52        "//tensorflow/compiler/xla/service:hlo_proto_cc",
53        ":tpu_driver_proto_cc",
54    ] + external_deps(),
55)
56
57cc_library(
58    name = "grpc_tpu_driver",
59    srcs = [
60        "grpc_tpu_driver.cc",
61    ],
62    hdrs = ["grpc_tpu_driver.h"],
63    deps = [
64        ":tpu_driver",
65        "//tensorflow/core/platform:logging",
66        "//tensorflow/compiler/xla:status",
67        "//tensorflow/compiler/xla:util",
68        "//tensorflow/compiler/xla:xla_data_proto_cc",
69        "//tensorflow/compiler/xla/service:hlo_proto_cc",
70        ":tpu_service_proto_cc",
71        ":tpu_driver_proto_cc",
72        tf_grpc_cc_dependency(),
73    ] + external_deps(),
74    alwayslink = 1,
75)
76
77cc_library(
78    name = "direct_tpu_driver",
79    srcs = ["direct_tpu_driver.cc"],
80    compatible_with = [],
81    deps = [
82        ":tpu_driver",
83        "@com_google_absl//absl/strings:str_format",
84        "//tensorflow/compiler/xla:statusor",
85        "//tensorflow/core/platform:logging",
86        "//tensorflow/compiler/xla:status",
87        "//tensorflow/compiler/xla:util",
88        "//tensorflow/compiler/xla:xla_data_proto_cc",
89        "//tensorflow/compiler/xla/service:hlo_proto_cc",
90        ":tpu_service_proto_cc",
91        ":tpu_driver_proto_cc",
92        "//tensorflow/compiler/xla/python/tpu_driver/client:libtpu",
93    ] + external_deps(),
94    alwayslink = 1,
95)
96
97cc_library(
98    name = "recording_tpu_driver",
99    srcs = [
100        "recording_tpu_driver.cc",
101    ],
102    deps = [
103        ":tpu_driver",
104        "@com_google_absl//absl/base",
105        "@com_google_absl//absl/types:optional",
106        "//tensorflow/core/platform:env",
107        "//tensorflow/core/platform:logging",
108        "//tensorflow/core/platform:stringpiece",
109        "//tensorflow/compiler/xla:status",
110        "//tensorflow/compiler/xla:util",
111        "//tensorflow/compiler/xla:xla_data_proto_cc",
112        ":tpu_driver_proto_cc",
113        ":tpu_service_proto_cc",
114        "//tensorflow/compiler/xla/service:hlo_proto_cc",
115    ] + external_deps(),
116    alwayslink = 1,
117)
118
119cc_library(
120    name = "pod_tpu_driver",
121    srcs = ["pod_tpu_driver.cc"],
122    deps = [
123        ":grpc_tpu_driver",
124        ":tpu_driver",
125        ":tpu_driver_proto_cc",
126        "@com_google_absl//absl/container:btree",
127        "@com_google_absl//absl/container:flat_hash_set",
128        "//tensorflow/compiler/xla/pjrt:semaphore",
129        "//tensorflow/compiler/xla/pjrt:worker_thread",
130        "//tensorflow/core:lib",
131        "//tensorflow/core:protos_all_cc",
132        tf_grpc_cc_dependency(),
133    ] + external_deps(),
134    alwayslink = 1,
135)
136
137go_grpc_library(
138    name = "tpu_service_go_grpc",
139    srcs = [":tpu_service_proto"],
140    compatible_with = ["//buildenv/target:gce"],
141    deps = [":tpu_service_go_proto"],
142)
143