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

Change not visible to public. 

PiperOrigin-RevId: 314878442
Change-Id: Ife98e7f85a9c800caa1c24153f5edac15ef40b7c
This commit is contained in:
Frank Chen 2020-06-05 00:15:14 -07:00 committed by TensorFlower Gardener
parent 6109ccc4ea
commit 3dcbf3c4fd
2 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/core/tpu/tpu_config_c_api.h
View File

@ -24,9 +24,13 @@ limitations under the License.
typedef struct TpuSerializedProto TpuSerializedProto; typedef struct TpuSerializedProto TpuSerializedProto;
extern "C" { namespace tensorflow {
namespace tpu {
class TpuMeshStateInterface;
} // namespace tpu
} // namespace tensorflow
TFTPU_CAPI_EXPORT bool TPUHostInitialized(); extern "C" {
TFTPU_CAPI_EXPORT void ConfigureDistributedTpuOp_DoWork( TFTPU_CAPI_EXPORT void ConfigureDistributedTpuOp_DoWork(
const size_t num_cores_per_host_size, const int32_t* num_cores_per_host, const size_t num_cores_per_host_size, const int32_t* num_cores_per_host,
@ -36,6 +40,7 @@ TFTPU_CAPI_EXPORT void ConfigureDistributedTpuOp_DoWork(
TFTPU_CAPI_EXPORT void WaitForDistributedTpuOp_DoWork( TFTPU_CAPI_EXPORT void WaitForDistributedTpuOp_DoWork(
const size_t num_hosts, const size_t num_cores_per_host, const size_t num_hosts, const size_t num_cores_per_host,
const int32_t** host_ordinal_to_global_core_id_map, const int32_t** host_ordinal_to_global_core_id_map,
tensorflow::tpu::TpuMeshStateInterface* tpu_mesh_state_interface,
size_t* tpu_topology_output_size, char** tpu_topology_output, size_t* tpu_topology_output_size, char** tpu_topology_output,
TF_Status* status); TF_Status* status);
@ -58,7 +63,6 @@ TFTPU_CAPI_EXPORT void TpuConfigurationApi_FreeInt32Array(int32_t* output);
} }
struct TfTpu_ConfigApiFn { struct TfTpu_ConfigApiFn {
TFTPU_ADD_FN_IN_STRUCT(TPUHostInitialized);
TFTPU_ADD_FN_IN_STRUCT(ConfigureDistributedTpuOp_DoWork); TFTPU_ADD_FN_IN_STRUCT(ConfigureDistributedTpuOp_DoWork);
TFTPU_ADD_FN_IN_STRUCT(WaitForDistributedTpuOp_DoWork); TFTPU_ADD_FN_IN_STRUCT(WaitForDistributedTpuOp_DoWork);
TFTPU_ADD_FN_IN_STRUCT(ShutdownDistributedTpuOp_DoWork); TFTPU_ADD_FN_IN_STRUCT(ShutdownDistributedTpuOp_DoWork);

1
tensorflow/core/tpu/tpu_library_loader.cc
View File

@ -41,7 +41,6 @@ Status SetTpuInitializeStructFns(void* library_handle) {
Status SetTpuConfigStructFns(void* library_handle) { Status SetTpuConfigStructFns(void* library_handle) {
auto* config_fn = ConfigApiFn(); auto* config_fn = ConfigApiFn();
TFTPU_SET_FN(config_fn, TPUHostInitialized);
TFTPU_SET_FN(config_fn, ConfigureDistributedTpuOp_DoWork); TFTPU_SET_FN(config_fn, ConfigureDistributedTpuOp_DoWork);
TFTPU_SET_FN(config_fn, WaitForDistributedTpuOp_DoWork); TFTPU_SET_FN(config_fn, WaitForDistributedTpuOp_DoWork);
TFTPU_SET_FN(config_fn, ShutdownDistributedTpuOp_DoWork); TFTPU_SET_FN(config_fn, ShutdownDistributedTpuOp_DoWork);