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

[TPU] Break dependencies between TPU stream executor and shared library loader

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: 353945938
Change-Id: I6c7a195467720489608ee236e2cbb2ba4d0d1444
This commit is contained in:
Frank Chen 2021-01-26 14:06:48 -08:00 committed by TensorFlower Gardener
parent 9ee7896d22
commit a6a4f88e85
9 changed files with 24 additions and 21 deletions

15
tensorflow/compiler/jit/BUILD
View File

@ -67,7 +67,10 @@ cc_library(
] + if_cuda_or_rocm([
":xla_gpu_device",
":xla_gpu_jit",
]) + if_with_tpu_support([":xla_tpu_device"]),
]) + if_with_tpu_support([
":xla_tpu_device",
":xla_tpu_jit",
]),
alwayslink = 1,
)
@ -101,6 +104,16 @@ cc_library(
alwayslink = 1,
)
cc_library(
name = "xla_tpu_jit",
visibility = ["//visibility:public"],
deps = if_libtpu([
"//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration",
"//tensorflow/stream_executor/tpu:tpu_transfer_manager",
]),
alwayslink = 1,
)
cc_library(
name = "xla_cpu_device",
srcs = ["xla_cpu_device.cc"],

7
tensorflow/core/tpu/BUILD
View File

@ -146,9 +146,7 @@ cc_library(
":tpu_executor_init_fns",
":tpu_library_init_fns",
":tpu_ops_c_api_hdrs",
"//tensorflow/compiler/xla/service:tpu_computation_placer",
"//tensorflow/core:lib",
"//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration",
"//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
],
)
@ -264,8 +262,8 @@ cc_library(
"//tensorflow/stream_executor/tpu:proto_helper",
"//tensorflow/stream_executor/tpu:status_helper",
"//tensorflow/stream_executor/tpu:tpu_executable_interface",
"//tensorflow/stream_executor/tpu:tpu_executor",
"//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
"//tensorflow/stream_executor/tpu:tpu_executor_hdrs",
"//tensorflow/stream_executor/tpu:tpu_platform_id",
"@com_google_absl//absl/types:span",
],
@ -279,10 +277,7 @@ cc_library(
deps = [
":tpu_api_dlsym_initializer",
"//tensorflow/core/tpu:tpu_on_demand_compiler",
"//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration",
"//tensorflow/core/tpu/ops",
"//tensorflow/stream_executor/tpu:tpu_executor",
"//tensorflow/stream_executor/tpu:tpu_transfer_manager",
],
)

1
tensorflow/core/tpu/graph_rewrite/BUILD
View File

@ -4,6 +4,7 @@ load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
default_visibility = [
"//tensorflow/compiler:__subpackages__",
"//tensorflow/core/tpu:__subpackages__",
"//tensorflow/stream_executor/tpu:__subpackages__",
],

2
tensorflow/core/tpu/kernels/BUILD
View File

@ -639,8 +639,8 @@ cc_library(
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core/tpu:tpu_ops_c_api_hdrs",
"//tensorflow/stream_executor/tpu:tpu_executor",
"//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
"//tensorflow/stream_executor/tpu:tpu_executor_hdrs",
"@com_google_absl//absl/types:variant",
],
alwayslink = 1,

5
tensorflow/core/tpu/tpu_api_dlsym_initializer.cc
View File

@ -20,11 +20,10 @@ limitations under the License.
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/tpu_api_dlsym_set_fn.h"
#if !defined(PLATFORM_GOOGLE)
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/core/tpu/tpu_initializer_helper.h"
#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_platform.h"
#endif
@ -56,8 +55,6 @@ Status InitializeTpuLibrary(void* library_handle) {
dlsym(library_handle, "TfTpu_Initialize"));
(*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1,
/*argv=*/argv_ptr.data());
RegisterTpuPlatform();
}
return s;

7
tensorflow/core/tpu/tpu_executor_dlsym_initializer.cc
View File

@ -21,11 +21,10 @@ limitations under the License.
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/tpu_api_dlsym_set_fn.h"
#if !defined(PLATFORM_GOOGLE)
#include "tensorflow/core/tpu/tpu_executor_api.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/core/tpu/tpu_initializer_helper.h"
#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_platform.h"
#endif
namespace tensorflow {
@ -53,8 +52,6 @@ Status InitializeTpuLibrary(void* library_handle) {
dlsym(library_handle, "TfTpu_Initialize"));
(*initialize_fn)(/*init_library=*/true, /*argc=*/argv_ptr.size() - 1,
/*argv=*/argv_ptr.data());
RegisterTpuPlatform();
}
return s;

3
tensorflow/stream_executor/tpu/BUILD
View File

@ -22,7 +22,6 @@ cc_library(
],
deps = [
"//tensorflow/c:tf_attrtype",
"//tensorflow/c:tf_status",
"//tensorflow/core/tpu:libtftpu_header",
"//tensorflow/stream_executor:stream_executor_headers",
],
@ -159,7 +158,7 @@ cc_library(
":tpu_topology_external",
"//tensorflow/core:lib",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor",
"//tensorflow/stream_executor:stream_executor_headers",
"//tensorflow/stream_executor/lib",
"@com_google_absl//absl/container:flat_hash_map",
],

4
tensorflow/stream_executor/tpu/c_api_decl.h
View File

@ -20,11 +20,13 @@ limitations under the License.
#include <stdint.h>
#include "tensorflow/c/tf_attrtype.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/tpu/libtftpu.h"
extern "C" {
struct TF_Status;
typedef struct TF_Status TF_Status;
// Maximum number of array elements to inline into structs for performance.
#define TPU_C_API_MAX_INLINED 6

1
tensorflow/stream_executor/tpu/tpu_executor_c_api.h
View File

@ -20,7 +20,6 @@ limitations under the License.
#include <stdint.h>
#include "tensorflow/c/tf_attrtype.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/tpu/libtftpu.h"
#include "tensorflow/stream_executor/tpu/c_api_decl.h"