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

Rename tpu_load_library to tpu_api_dlsym_initializer. 

This better reflects what the load library module does.

PiperOrigin-RevId: 317763973
Change-Id: I66169394c2f1d77b260323cb81d76f0f4e167579
This commit is contained in:
Frank Chen 2020-06-22 16:58:50 -07:00 committed by TensorFlower Gardener
parent 6b0b4876f8
commit 23e387975e
18 changed files with 141 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/BUILD
View File

@ -2259,7 +2259,7 @@ tf_cuda_library(
"//tensorflow/core/platform/default/build_config:platformlib",
"//tensorflow/core/profiler/lib:annotated_traceme",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api_dlsym_initializer",
"//tensorflow/core/util:einsum_op_util",
"//tensorflow/core/util:padding",
"//tensorflow/core/util:port",

2
tensorflow/core/framework/load_library.cc
View File

@ -22,7 +22,7 @@ limitations under the License.
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mem.h"
#if !defined(IS_MOBILE_PLATFORM)
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api_dlsym_initializer.h"
#endif // IS_MOBILE_PLATFORM
namespace tensorflow {

24
tensorflow/core/tpu/BUILD
View File

@ -107,15 +107,31 @@ cc_library(
)
cc_library(
name = "tpu_library_loader",
name = "tpu_api",
srcs = ["tpu_api.cc"],
hdrs = ["tpu_api.h"],
deps = [
":libtftpu_header",
":tpu_config_c_api",
"//tensorflow/core/tpu/kernels:tpu_compile_c_api_hdrs",
"//tensorflow/core/tpu/kernels:tpu_mesh_state_c_api_hdrs",
"//tensorflow/core/tpu/kernels:tpu_util_c_api_hdrs",
"//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
"//tensorflow/stream_executor/tpu:tpu_node_context_c_api_hdrs",
],
)
cc_library(
name = "tpu_api_dlsym_initializer",
srcs = if_windows(
["tpu_library_loader_windows.cc"],
otherwise = ["tpu_library_loader.cc"],
["tpu_api_dlsym_initializer_windows.cc"],
otherwise = ["tpu_api_dlsym_initializer.cc"],
),
hdrs = ["tpu_library_loader.h"],
hdrs = ["tpu_api_dlsym_initializer.h"],
visibility = ["//visibility:public"],
deps = [
":libtftpu_header",
":tpu_api",
":tpu_config_c_api",
":tpu_library_init_fns",
"//tensorflow/core/platform:errors",

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

@ -72,10 +72,10 @@ tf_kernel_library(
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:refcount",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/core/tpu:tpu_config_c_api",
"//tensorflow/core/tpu:tpu_configuration",
"//tensorflow/core/tpu:tpu_defs",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/stream_executor/tpu:proto_helper",
],
alwayslink = 1,
@ -224,7 +224,7 @@ cc_library(
"//tensorflow/compiler/xla/service",
"//tensorflow/core:framework",
"//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
],
)

2
tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
View File

@ -24,10 +24,10 @@ limitations under the License.
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/core/tpu/tpu_config_c_api.h"
#include "tensorflow/core/tpu/tpu_configuration.h"
#include "tensorflow/core/tpu/tpu_defs.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/stream_executor/tpu/proto_helper.h"
namespace tensorflow {

2
tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
View File

@ -21,7 +21,7 @@ limitations under the License.
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
namespace tensorflow {

57
tensorflow/core/tpu/tpu_api.cc Normal file
View File

@ -0,0 +1,57 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/tpu/tpu_api.h"
namespace tensorflow {
namespace tpu {
TfTpu_BaseFn* InitializeApiFn() {
static TfTpu_BaseFn base_fn;
return &base_fn;
}
TfTpu_ConfigApiFn* ConfigApiFn() {
static TfTpu_ConfigApiFn config_api_fn;
return &config_api_fn;
}
TfTpu_MeshStateApiFn* MeshStateApiFn() {
static TfTpu_MeshStateApiFn mesh_state_api_fn;
return &mesh_state_api_fn;
}
TfTpu_CompileApiFn* CompileApiFn() {
static TfTpu_CompileApiFn compile_api_fn;
return &compile_api_fn;
}
TfTpu_ExecutorApiFn* ExecutorApiFn() {
static TfTpu_ExecutorApiFn executor_api_fn;
return &executor_api_fn;
}
TfTpu_NodeContextApiFn* NodeContextApiFn() {
static TfTpu_NodeContextApiFn node_context_api_fn;
return &node_context_api_fn;
}
TfTpu_UtilApiFn* UtilApiFn() {
static TfTpu_UtilApiFn util_api_fn;
return &util_api_fn;
}
} // namespace tpu
} // namespace tensorflow

12
tensorflow/core/tpu/tpu_library_loader.h → tensorflow/core/tpu/tpu_api.h
View File

@ -13,10 +13,9 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
#define TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
#ifndef TENSORFLOW_CORE_TPU_TPU_API_H_
#define TENSORFLOW_CORE_TPU_TPU_API_H_
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_util_c_api.h"
@ -25,13 +24,9 @@ limitations under the License.
#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
// LINT.IfChange
namespace tensorflow {
namespace tpu {
Status InitializeTpuLibrary(void* library_handle);
// TODO(frankchn): Separate out API functions from the loader.
TfTpu_BaseFn* InitializeApiFn();
TfTpu_ConfigApiFn* ConfigApiFn();
@ -48,6 +43,5 @@ TfTpu_UtilApiFn* UtilApiFn();
} // namespace tpu
} // namespace tensorflow
// LINT.ThenChange(//tensorflow/core/tpu/tpu_library_loader_windows.cc)
#endif // TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
#endif // TENSORFLOW_CORE_TPU_TPU_API_H_

40
tensorflow/core/tpu/tpu_library_loader.cc → tensorflow/core/tpu/tpu_api_dlsym_initializer.cc
View File

@ -13,14 +13,13 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// TODO(frankchn): Rename to `tpu_api_dlsym_initializer` or similar.
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api_dlsym_initializer.h"
#include <dlfcn.h>
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_platform.h"
@ -39,41 +38,6 @@ namespace tpu {
#include "tensorflow/core/tpu/tpu_library_init_fns.inc"
TfTpu_BaseFn* InitializeApiFn() {
static TfTpu_BaseFn base_fn;
return &base_fn;
}
TfTpu_ConfigApiFn* ConfigApiFn() {
static TfTpu_ConfigApiFn config_api_fn;
return &config_api_fn;
}
TfTpu_MeshStateApiFn* MeshStateApiFn() {
static TfTpu_MeshStateApiFn mesh_state_api_fn;
return &mesh_state_api_fn;
}
TfTpu_CompileApiFn* CompileApiFn() {
static TfTpu_CompileApiFn compile_api_fn;
return &compile_api_fn;
}
TfTpu_ExecutorApiFn* ExecutorApiFn() {
static TfTpu_ExecutorApiFn executor_api_fn;
return &executor_api_fn;
}
TfTpu_NodeContextApiFn* NodeContextApiFn() {
static TfTpu_NodeContextApiFn node_context_api_fn;
return &node_context_api_fn;
}
TfTpu_UtilApiFn* UtilApiFn() {
static TfTpu_UtilApiFn util_api_fn;
return &util_api_fn;
}
Status InitializeTpuLibrary(void* library_handle) {
bool shared_object_loaded = true;
if (library_handle == nullptr) {

38
tensorflow/core/tpu/tpu_api_dlsym_initializer.h Normal file
View File

@ -0,0 +1,38 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_TPU_TPU_API_DLSYM_INITIALIZER_H_
#define TENSORFLOW_CORE_TPU_TPU_API_DLSYM_INITIALIZER_H_
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_util_c_api.h"
#include "tensorflow/core/tpu/libtftpu.h"
#include "tensorflow/core/tpu/tpu_config_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
// LINT.IfChange
namespace tensorflow {
namespace tpu {
Status InitializeTpuLibrary(void* library_handle);
} // namespace tpu
} // namespace tensorflow
// LINT.ThenChange(//tensorflow/core/tpu/tpu_api_dlsym_initializer_windows.cc)
#endif // TENSORFLOW_CORE_TPU_TPU_API_DLSYM_INITIALIZER_H_

20
tensorflow/core/tpu/tpu_library_loader_windows.cc → tensorflow/core/tpu/tpu_api_dlsym_initializer_windows.cc
View File

@ -15,28 +15,14 @@ limitations under the License.
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api_dlsym_initializer.h"
// Reminder: Update tpu_library_loader.cc if you are adding new publicly
// visible methods.
// Reminder: Update tpu_api_dlsym_initializer_windows.cc if you are adding new
// publicly visible methods.
namespace tensorflow {
namespace tpu {
TfTpu_BaseFn* InitializeApiFn() { return nullptr; }
TfTpu_ConfigApiFn* ConfigApiFn() { return nullptr; }
TfTpu_MeshStateApiFn* MeshStateApiFn() { return nullptr; }
TfTpu_CompileApiFn* CompileApiFn() { return nullptr; }
TfTpu_ExecutorApiFn* ExecutorApiFn() { return nullptr; }
TfTpu_NodeContextApiFn* NodeContextApiFn() { return nullptr; }
TfTpu_UtilApiFn* UtilApiFn() { return nullptr; }
Status InitializeTpuLibrary(void* library_handle) {
return errors::Unimplemented(
"Loading TPU library is not supported on Windows.");

12
tensorflow/stream_executor/tpu/BUILD
View File

@ -70,7 +70,7 @@ cc_library(
":status_helper",
":tpu_executor_c_api_hdrs",
":tpu_stream_interface",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor:stream",
],
)
@ -81,7 +81,7 @@ cc_library(
deps = [
":tpu_executor_c_api_hdrs",
"//tensorflow/core/platform:types",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor:stream",
],
)
@ -101,7 +101,7 @@ cc_library(
":tpu_timer",
"//tensorflow/c:tf_status",
"//tensorflow/core:lib",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor:stream",
"//tensorflow/stream_executor/lib",
"@com_google_absl//absl/container:flat_hash_map",
@ -151,7 +151,7 @@ cc_library(
"//tensorflow/compiler/xla/service:stream_pool",
"//tensorflow/compiler/xla/service:transfer_manager",
"//tensorflow/core:framework",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor:device_memory_allocator",
"//tensorflow/stream_executor/lib",
"@com_google_absl//absl/memory",
@ -169,7 +169,7 @@ cc_library(
":tpu_platform_interface",
"//tensorflow/c:tf_status",
"//tensorflow/core/platform:types",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor:stream",
"@com_google_absl//absl/container:flat_hash_map",
],
@ -201,7 +201,7 @@ cc_library(
"//tensorflow/compiler/xla:xla_data_proto_cc",
"//tensorflow/compiler/xla/service:shaped_buffer",
"//tensorflow/compiler/xla/service:transfer_manager",
"//tensorflow/core/tpu:tpu_library_loader",
"//tensorflow/core/tpu:tpu_api",
"//tensorflow/stream_executor:stream",
],
)

2
tensorflow/stream_executor/tpu/tpu_executor.cc
View File

@ -17,7 +17,7 @@ limitations under the License.
#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/device_memory.h"
#include "tensorflow/stream_executor/lib/status.h"
#include "tensorflow/stream_executor/tpu/c_api_conversions.h"

2
tensorflow/stream_executor/tpu/tpu_node_context.cc
View File

@ -17,7 +17,7 @@ limitations under the License.
#include "tensorflow/compiler/xla/service/backend.h"
#include "tensorflow/compiler/xla/service/platform_util.h"
#include "tensorflow/compiler/xla/service/transfer_manager.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/device_memory_allocator.h"
#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"

2
tensorflow/stream_executor/tpu/tpu_platform.cc
View File

@ -16,7 +16,7 @@ limitations under the License.
#include "tensorflow/stream_executor/tpu/tpu_platform.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/platform.h"
#include "tensorflow/stream_executor/tpu/status_helper.h"
#include "tensorflow/stream_executor/tpu/tpu_executor.h"

2
tensorflow/stream_executor/tpu/tpu_stream.h
View File

@ -16,7 +16,7 @@ limitations under the License.
#ifndef TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_STREAM_H_
#define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_STREAM_H_
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/stream_executor_internal.h"
#include "tensorflow/stream_executor/tpu/c_api_conversions.h"
#include "tensorflow/stream_executor/tpu/status_helper.h"

2
tensorflow/stream_executor/tpu/tpu_timer.h
View File

@ -17,7 +17,7 @@ limitations under the License.
#define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_TIMER_H_
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/stream_executor_internal.h"
#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"

2
tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
View File

@ -17,7 +17,7 @@ limitations under the License.
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/tpu/tpu_library_loader.h"
#include "tensorflow/core/tpu/tpu_api.h"
#include "tensorflow/stream_executor/device_memory.h"
#include "tensorflow/stream_executor/tpu/c_api_conversions.h"
#include "tensorflow/stream_executor/tpu/proto_helper.h"