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_helpers",
":xla_op_registry", ":xla_op_registry",
":xla_resource", ":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:common",
"//tensorflow/compiler/jit:flags", "//tensorflow/compiler/jit:flags",
"//tensorflow/compiler/jit:shape_inference", "//tensorflow/compiler/jit:shape_inference",
"//tensorflow/compiler/mlir:array_container_utils", "//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/mlir/tensorflow:compile_mlir_util_no_tf_dialect_passes",
"//tensorflow/compiler/xla:protobuf_util", "//tensorflow/compiler/xla:protobuf_util",
"//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:shape_util",
@ -373,12 +378,9 @@ cc_library(
"//tensorflow/core:lib_internal", "//tensorflow/core:lib_internal",
"//tensorflow/core:ops", "//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc", "//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/algorithm:container", ] + if_libtpu([
"@com_google_absl//absl/container:flat_hash_map", ":xla_tpu_backend_registration",
"@com_google_absl//absl/memory", ]),
"@com_google_absl//absl/types:span",
"@com_google_absl//absl/types:variant",
],
alwayslink = 1, alwayslink = 1,
) )
@ -405,6 +407,18 @@ cc_library(
alwayslink = 1, 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( cc_library(
name = "xla_context", name = "xla_context",
srcs = [ 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( cc_library(
name = "tpu_node_device_util", name = "tpu_node_device_util",
srcs = ["tpu_node_device_util.cc"], srcs = ["tpu_node_device_util.cc"],
@ -149,7 +136,6 @@ cc_library(
":libtftpu_header", ":libtftpu_header",
":tpu_api", ":tpu_api",
":tpu_api_dlsym_set_fn", ":tpu_api_dlsym_set_fn",
":tpu_compilation_device",
":tpu_executor_init_fns", ":tpu_executor_init_fns",
":tpu_library_init_fns", ":tpu_library_init_fns",
":tpu_ops_c_api_hdrs", ":tpu_ops_c_api_hdrs",
@ -285,7 +271,6 @@ cc_library(
visibility = ["//visibility:public"], visibility = ["//visibility:public"],
deps = [ deps = [
":tpu_api_dlsym_initializer", ":tpu_api_dlsym_initializer",
":tpu_compilation_device",
"//tensorflow/core/tpu:tpu_on_demand_compiler", "//tensorflow/core/tpu:tpu_on_demand_compiler",
"//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration", "//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration",
"//tensorflow/core/tpu/ops", "//tensorflow/core/tpu/ops",