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

This is an internal change not visible to the public.

Browse Source 
PiperOrigin-RevId: 311661259
Change-Id: I8443f6037e3cabca1dca72ce1748eea8fd71770c
This commit is contained in:
Frank Chen 2020-05-14 20:45:52 -07:00 committed by TensorFlower Gardener
parent a9c0ce87a6
commit 97f2fffe7e
1 changed files with 19 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -17,6 +17,7 @@ limitations under the License.
#define TENSORFLOW_CORE_TPU_TPU_CONFIG_C_API_H_ #define TENSORFLOW_CORE_TPU_TPU_CONFIG_C_API_H_
#include <cstddef> #include <cstddef>
#include <cstdint>
#include "tensorflow/c/tf_status.h" #include "tensorflow/c/tf_status.h"
@ -26,29 +27,33 @@ extern "C" {
bool TPUHostInitialized(); bool TPUHostInitialized();
// TODO(frankchn): Modify API to take in raw values instead of Tensors. void ConfigureDistributedTpuOp_DoWork(const size_t num_cores_per_host_size,
void ConfigureDistributedTpuOp_DoWork(size_t input_size, const int32_t* num_cores_per_host,
TpuSerializedProto** inputs, size_t* host_config_output_size,
TpuSerializedProto* output, char** host_config_output,
TF_Status* status); TF_Status* status);
void WaitForDistributedTpuOp_DoWork(size_t input_size, void WaitForDistributedTpuOp_DoWork(
TpuSerializedProto** inputs, const size_t num_hosts, const size_t num_cores_per_host,
TpuSerializedProto* output, const int32_t** host_ordinal_to_global_core_id_map,
TF_Status* status); size_t* tpu_topology_output_size, char** tpu_topology_output,
TF_Status* status);
void ShutdownDistributedTpuOp_DoWork(TF_Status* status); void ShutdownDistributedTpuOp_DoWork(TF_Status* status);
void InitializeHostForDistributedTpuOp_DoWork( void InitializeHostForDistributedTpuOp_DoWork(
size_t input_size, TpuSerializedProto** inputs, const size_t tpu_host_config_size, const char* tpu_host_config,
bool enable_whole_mesh_compilations, TpuSerializedProto* output, const bool enable_whole_mesh_compilations, size_t* core_id_output_size,
TF_Status* status); int32_t** core_id_output, TF_Status* status);
void SetGlobalTPUArrayOp_DoWork(size_t input_size, TpuSerializedProto** inputs, void SetGlobalTPUArrayOp_DoWork(const size_t tpu_topology_size,
TF_Status* status); const char* tpu_topology, TF_Status* status);
void DisconnectDistributedTpuChipsOp_DoWork(TpuSerializedProto* output, void DisconnectDistributedTpuChipsOp_DoWork(int32_t* number_of_chips_output,
TF_Status* status); TF_Status* status);
void TpuConfigurationApi_FreeCharArray(char* output);
void TpuConfigurationApi_FreeInt32Array(int32_t* output);
} }
#endif // TENSORFLOW_CORE_TPU_TPU_CONFIG_C_API_H_ #endif // TENSORFLOW_CORE_TPU_TPU_CONFIG_C_API_H_