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

Add APIs for Buffer Linearization Support in TPUs

Browse Source 
PiperOrigin-RevId: 325525811
Change-Id: I964daa0ba93bba5707c4714ac8de3564bba58c06
This commit is contained in:
Frank Chen 2020-08-07 15:58:52 -07:00 committed by TensorFlower Gardener
parent f6cfee3dff
commit 0d2f665129
2 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/tpu/tpu_library_init_fns.inc
View File

@ -161,6 +161,8 @@ tensorflow::Status SetExecutorStructFn(void* library_handle) {
TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromDevice);
TFTPU_SET_FN(executor_fn, TpuTransferManager_GetByteSizeRequirement);
TFTPU_SET_FN(executor_fn, TpuTransferManager_WriteSingleTupleIndexTable);
TFTPU_SET_FN(executor_fn, TpuTransferManager_LinearizeToBuffers);
TFTPU_SET_FN(executor_fn, TpuTransferManager_FreeBuffers);
TFTPU_SET_FN(executor_fn, TpuComputationPlacer_New);
TFTPU_SET_FN(executor_fn, TpuComputationPlacer_Free);

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

@ -182,6 +182,11 @@ void TpuTransferManager_WriteSingleTupleIndexTable(
XLA_TransferManager* manager, SE_Stream* stream,
SE_DeviceMemoryBase* elements, size_t elements_len, XLA_Shape* shape,
SE_DeviceMemoryBase* region, SE_Status* status);
void TpuTransferManager_LinearizeToBuffers(
XLA_TransferManager* manager, XLA_Literal* c_literal, char*** buffers_array,
int64_t** buffers_size, int64_t* buffers_array_size, SE_Status* status);
void TpuTransferManager_FreeBuffers(char** buffers_array, int64_t* buffers_size,
int64_t buffers_array_size);
XLA_ComputationPlacer* TpuComputationPlacer_New();
void TpuComputationPlacer_Free(XLA_ComputationPlacer* placer);
@ -336,6 +341,8 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_TransferLiteralFromDevice);
TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_GetByteSizeRequirement);
TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_WriteSingleTupleIndexTable);
TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_LinearizeToBuffers);
TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_FreeBuffers);
TFTPU_ADD_FN_IN_STRUCT(TpuComputationPlacer_New);
TFTPU_ADD_FN_IN_STRUCT(TpuComputationPlacer_Free);