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