experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[TPU] Colocate tpu_compilation_device registration with other registration in compiler/tf2xla

Browse Source 
This is part of a series of changes to move TPU-related code to better locations so that the TensorFlow build isn't confused and TPU-based TF can be built without the define=framework_shared_object=false flag.

PiperOrigin-RevId: 353341905
Change-Id: I7cd820aab37c3d4f1a838967a87dc462491e446f
This commit is contained in:
Frank Chen 2021-01-22 16:32:02 -08:00 committed by TensorFlower Gardener
parent 8d33fd7426
commit bdc490c868
3 changed files with 21 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
tensorflow/compiler/tf2xla/BUILD
View File

@ -350,11 +350,16 @@ cc_library(
":xla_helpers",
":xla_op_registry",
":xla_resource",
"//tensorflow/compiler/mlir:mlir_bridge_rollout_policy",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/types:span",
"@com_google_absl//absl/types:variant",
"//tensorflow/compiler/jit:common",
"//tensorflow/compiler/jit:flags",
"//tensorflow/compiler/jit:shape_inference",
"//tensorflow/compiler/mlir:array_container_utils",
"//tensorflow/compiler/mlir:mlir_bridge_rollout_policy",
"//tensorflow/compiler/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
"//tensorflow/compiler/xla:protobuf_util",
"//tensorflow/compiler/xla:shape_util",
@ -373,12 +378,9 @@ cc_library(
"//tensorflow/core:lib_internal",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/types:span",
"@com_google_absl//absl/types:variant",
],
] + if_libtpu([
":xla_tpu_backend_registration",
]),
alwayslink = 1,
)
@ -405,6 +407,18 @@ cc_library(
alwayslink = 1,
)
cc_library(
name = "xla_tpu_backend_registration",
srcs = ["xla_tpu_backend.cc"],
visibility = ["//visibility:public"],
deps = [
":xla_op_registry",
"//tensorflow/core/tpu:tpu_defs",
"//tensorflow/core/tpu:tpu_node_device_util",
],
alwayslink = 1,
)
cc_library(
name = "xla_context",
srcs = [

0
tensorflow/core/tpu/tpu_compilation_device.cc → tensorflow/compiler/tf2xla/xla_tpu_backend.cc
View File

15
tensorflow/core/tpu/BUILD
View File

@ -55,19 +55,6 @@ cc_library(
],
)
cc_library(
name = "tpu_compilation_device",
srcs = ["tpu_compilation_device.cc"],
visibility = ["//visibility:public"],
deps = [
":tpu_defs",
":tpu_node_device_util",
"//tensorflow/compiler/tf2xla:xla_compiler",
"//tensorflow/compiler/tf2xla:xla_op_registry",
],
alwayslink = 1,
)
cc_library(
name = "tpu_node_device_util",
srcs = ["tpu_node_device_util.cc"],
@ -149,7 +136,6 @@ cc_library(
":libtftpu_header",
":tpu_api",
":tpu_api_dlsym_set_fn",
":tpu_compilation_device",
":tpu_executor_init_fns",
":tpu_library_init_fns",
":tpu_ops_c_api_hdrs",
@ -285,7 +271,6 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
":tpu_api_dlsym_initializer",
":tpu_compilation_device",
"//tensorflow/core/tpu:tpu_on_demand_compiler",
"//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration",
"//tensorflow/core/tpu/ops",