• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
2load("//tensorflow:tensorflow.bzl", "tf_grpc_cc_dependencies")
3load(
4    "//tensorflow:tensorflow.bzl",
5    "if_libtpu",
6    "tf_cc_binary",
7    "tf_cc_test",
8)
9load(
10    "//tensorflow/core/platform:build_config.bzl",
11    "tf_proto_library",
12)
13load(
14    "//tensorflow/compiler/xla:xla.bzl",
15    "xla_py_grpc_library",
16)
17
18package(
19    default_visibility = ["//tensorflow:internal"],
20    licenses = ["notice"],
21)
22
23tf_proto_library(
24    name = "xla_service_proto",
25    srcs = ["xla_service.proto"],
26    has_services = 1,
27    cc_api_version = 2,
28    create_grpc_library = True,
29    create_service = True,
30    protodeps = [
31        "//tensorflow/compiler/xla:xla_proto",
32    ],
33    use_grpc_namespace = True,
34    visibility = ["//visibility:public"],
35)
36
37cc_library(
38    name = "grpc_stub",
39    srcs = ["grpc_stub.cc"],
40    hdrs = ["grpc_stub.h"],
41    deps = [
42        ":xla_service_proto_cc",
43        "//tensorflow/compiler/xla:service_interface",
44        "//tensorflow/compiler/xla:xla_data_cc_grpc_proto",
45        "//tensorflow/compiler/xla:xla_data_proto_cc",
46        "//tensorflow/compiler/xla/rpc:xla_service_cc_grpc_proto",
47        "//tensorflow/core/distributed_runtime/rpc:grpc_util",
48    ],
49)
50
51cc_library(
52    name = "grpc_service_main_library",
53    srcs = ["grpc_service_main.cc"],
54    deps = [
55        ":grpc_service",
56        "@com_google_absl//absl/strings:str_format",
57        "//tensorflow/compiler/xla/service:platform_util",
58        "//tensorflow/core:framework_internal",
59        "//tensorflow/core:lib",
60    ] + tf_grpc_cc_dependencies() + if_libtpu(
61        if_false = ["//tensorflow/compiler/xla/service:cpu_plugin"],
62        if_true = [],
63    ),
64)
65
66tf_cc_binary(
67    name = "grpc_service_main_cpu",
68    deps = [
69        ":grpc_service_main_library",
70        "//tensorflow/compiler/xla/service:cpu_plugin",
71    ],
72)
73
74tf_cc_test(
75    name = "grpc_client_test",
76    srcs = ["grpc_client_test.cc"],
77    data = [
78        "//tensorflow/compiler/xla/rpc:grpc_service_main_cpu",
79    ],
80    deps = [
81        ":grpc_stub",
82        "@com_google_absl//absl/strings:str_format",
83        "//tensorflow/compiler/xla/client",
84        "//tensorflow/compiler/xla/client:xla_builder",
85        "//tensorflow/compiler/xla/tests:literal_test_util",
86        "//tensorflow/core:framework_internal",
87        "//tensorflow/core:lib",
88        "//tensorflow/core:test",
89        "//tensorflow/core:test_main",
90    ] + tf_grpc_cc_dependencies(),
91)
92
93cc_library(
94    name = "grpc_service",
95    srcs = ["grpc_service.cc"],
96    hdrs = ["grpc_service.h"],
97    deps = [
98        ":xla_service_proto_cc",
99        ":xla_service_cc_grpc_proto",
100        "//tensorflow/compiler/xla/service",
101        "//tensorflow/compiler/xla/service:platform_util",
102        "//tensorflow/core/distributed_runtime/rpc:grpc_util",
103    ] + tf_grpc_cc_dependencies(),
104)
105
106# copybara:uncomment_begin(google-only)
107# py_proto_library(
108#     name = "xla_py_pb2",
109#     has_services = True,
110#     api_version = 2,
111#     visibility = ["//visibility:public"],
112#     deps = [":xla_service_proto"],
113# )
114#
115# xla_py_grpc_library(
116#     name = "xla_py_pb2_grpc",
117#     srcs = [":xla_service_proto"],
118#     deps = [":xla_py_pb2"],
119# )
120# copybara:uncomment_end
121