• 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_dependencies")
4load(
5    "//tensorflow/compiler/xla/python/tpu_driver:platform/external/tools.bzl",
6    "external_deps",
7    "go_grpc_library",
8)
9
10licenses(["notice"])
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    create_grpc_library = True,
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        "//tensorflow/core/platform:logging",
47        "//tensorflow/compiler/xla:status",
48        "//tensorflow/compiler/xla:statusor",
49        "//tensorflow/compiler/xla:util",
50        "//tensorflow/compiler/xla:xla_data_proto_cc",
51        "//tensorflow/compiler/xla/service:hlo_proto_cc",
52        ":tpu_driver_proto_cc",
53    ] + external_deps(),
54)
55
56cc_library(
57    name = "grpc_tpu_driver",
58    srcs = [
59        "grpc_tpu_driver.cc",
60    ],
61    hdrs = ["grpc_tpu_driver.h"],
62    deps = [
63        ":tpu_driver",
64        "//tensorflow/core/platform:logging",
65        "//tensorflow/compiler/xla:status",
66        "//tensorflow/compiler/xla:util",
67        "//tensorflow/compiler/xla:xla_data_proto_cc",
68        "//tensorflow/compiler/xla/service:hlo_proto_cc",
69        ":tpu_service_proto_cc",
70        ":tpu_service_cc_grpc_proto",
71        ":tpu_driver_proto_cc",
72    ] + tf_grpc_cc_dependencies() + external_deps(),
73    alwayslink = 1,
74)
75
76cc_library(
77    name = "direct_tpu_driver",
78    srcs = ["direct_tpu_driver.cc"],
79    compatible_with = [],
80    deps = [
81        ":tpu_driver",
82        "@com_google_absl//absl/strings:str_format",
83        "//tensorflow/compiler/xla:statusor",
84        "//tensorflow/core/platform:logging",
85        "//tensorflow/compiler/xla:status",
86        "//tensorflow/compiler/xla:util",
87        "//tensorflow/compiler/xla:xla_data_proto_cc",
88        "//tensorflow/compiler/xla/service:hlo_proto_cc",
89        ":tpu_service_proto_cc",
90        ":tpu_driver_proto_cc",
91        "//tensorflow/compiler/xla/python/tpu_driver/client:libtpu",
92    ] + external_deps(),
93    alwayslink = 1,
94)
95
96cc_library(
97    name = "recording_tpu_driver",
98    srcs = [
99        "recording_tpu_driver.cc",
100    ],
101    deps = [
102        ":tpu_driver",
103        "@com_google_absl//absl/base",
104        "//tensorflow/core/platform:env",
105        "//tensorflow/core/platform:logging",
106        "//tensorflow/compiler/xla:status",
107        "//tensorflow/compiler/xla:util",
108        "//tensorflow/compiler/xla:xla_data_proto_cc",
109        ":tpu_driver_proto_cc",
110        ":tpu_service_proto_cc",
111        ":tpu_service_cc_grpc_proto",
112        "//tensorflow/compiler/xla/service:hlo_proto_cc",
113    ] + external_deps(),
114    alwayslink = 1,
115)
116
117cc_library(
118    name = "pod_tpu_driver",
119    srcs = ["pod_tpu_driver.cc"],
120    deps = [
121        ":grpc_tpu_driver",
122        ":tpu_driver",
123        ":tpu_driver_proto_cc",
124        "@com_google_absl//absl/container:btree",
125        "@com_google_absl//absl/container:flat_hash_set",
126        "//tensorflow/compiler/xla/pjrt:semaphore",
127        "//tensorflow/compiler/xla/pjrt:worker_thread",
128        "//tensorflow/core:lib",
129        "//tensorflow/core:protos_all_cc",
130    ] + tf_grpc_cc_dependencies() + external_deps(),
131    alwayslink = 1,
132)
133
134go_grpc_library(
135    name = "tpu_service_go_grpc",
136    srcs = [":tpu_service_proto"],
137    compatible_with = ["//buildenv/target:gce"],
138    deps = [":tpu_service_go_proto"],
139)
140