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

Add TpuExecutable_FreeXlaShapeIndexArray and TpuExecutable_FreeMaybeOwningDeviceMemoryArray to free appropriate memory

Browse Source 
PiperOrigin-RevId: 351234388
Change-Id: I9ac9c712ad3a2e715d1a849bce7fbc6b4c858c43
This commit is contained in:
Frank Chen 2021-01-11 14:19:30 -08:00 committed by TensorFlower Gardener
parent 756840185f
commit 986d13e168
3 changed files with 20 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -141,6 +141,8 @@ tensorflow::Status SetExecutorStructFn(void* library_handle) {
TFTPU_SET_FN(executor_fn, TpuCompiler_Compile); TFTPU_SET_FN(executor_fn, TpuCompiler_Compile);
TFTPU_SET_FN(executor_fn, TpuCompiler_ShapeSize); TFTPU_SET_FN(executor_fn, TpuCompiler_ShapeSize);
TFTPU_SET_FN(executor_fn, TpuExecutable_ExecuteAsyncOnStream); TFTPU_SET_FN(executor_fn, TpuExecutable_ExecuteAsyncOnStream);
TFTPU_SET_FN(executor_fn, TpuExecutable_FreeXlaShapeIndexArray);
TFTPU_SET_FN(executor_fn, TpuExecutable_FreeMaybeOwningDeviceMemoryArray);
TFTPU_SET_FN(executor_fn, TpuExecutable_Fingerprint); TFTPU_SET_FN(executor_fn, TpuExecutable_Fingerprint);
TFTPU_SET_FN(executor_fn, TpuExecutable_HloModule); TFTPU_SET_FN(executor_fn, TpuExecutable_HloModule);
TFTPU_SET_FN(executor_fn, TpuExecutable_Free); TFTPU_SET_FN(executor_fn, TpuExecutable_Free);

6
tensorflow/core/tpu/tpu_on_demand_compiler.cc
View File

@ -161,7 +161,8 @@ class TpuExecutable : public TpuExecutableInterface {
output.AddAliasedIndex( output.AddAliasedIndex(
ApiConverter::FromC(&se_execution_output.aliased_indices[i])); ApiConverter::FromC(&se_execution_output.aliased_indices[i]));
} }
ApiConverter::Free(se_execution_output.aliased_indices); ExecutorApiFn()->TpuExecutable_FreeXlaShapeIndexArrayFn(
se_execution_output.aliased_indices);
for (int i = 0; i < se_execution_output.to_be_released_size; ++i) { for (int i = 0; i < se_execution_output.to_be_released_size; ++i) {
output.AddToBeReleased( output.AddToBeReleased(
@ -170,7 +171,8 @@ class TpuExecutable : public TpuExecutableInterface {
.Release() .Release()
.value()); .value());
} }
delete[] se_execution_output.to_be_released; ExecutorApiFn()->TpuExecutable_FreeMaybeOwningDeviceMemoryArrayFn(
se_execution_output.to_be_released);
return output; return output;
} }

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

@ -309,6 +309,18 @@ TFTPU_CAPI_EXPORT void TpuExecutable_ExecuteAsyncOnStream(
SE_HloExecutionProfile* hlo_execution_profile, SE_HloExecutionProfile* hlo_execution_profile,
SE_ExecutionOutput* se_output, TF_Status* status); SE_ExecutionOutput* se_output, TF_Status* status);
// This frees the XLA_ShapeIndex* array allocated when se_output is returned by
// TpuExecutable_ExecuteAsyncOnStream.
TFTPU_CAPI_EXPORT void TpuExecutable_FreeXlaShapeIndexArray(
XLA_ShapeIndex* array);
// This frees the SE_MaybeOwningDeviceMemory* array allocated when se_output is
// returned by TpuExecutable_ExecuteAsyncOnStream.
// Note that this only frees the heap-allocated array itself, and does not
// free any of the underlying device memory.
TFTPU_CAPI_EXPORT void TpuExecutable_FreeMaybeOwningDeviceMemoryArray(
SE_MaybeOwningDeviceMemory* array);
TFTPU_CAPI_EXPORT void TpuExecutable_Fingerprint(SE_Executable* executable, TFTPU_CAPI_EXPORT void TpuExecutable_Fingerprint(SE_Executable* executable,
const char** fingerprint, const char** fingerprint,
size_t* size); size_t* size);
@ -468,6 +480,8 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuCompiler_Compile); TFTPU_ADD_FN_IN_STRUCT(TpuCompiler_Compile);
TFTPU_ADD_FN_IN_STRUCT(TpuCompiler_ShapeSize); TFTPU_ADD_FN_IN_STRUCT(TpuCompiler_ShapeSize);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_ExecuteAsyncOnStream); TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_ExecuteAsyncOnStream);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_FreeXlaShapeIndexArray);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_FreeMaybeOwningDeviceMemoryArray);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_Fingerprint); TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_Fingerprint);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_HloModule); TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_HloModule);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_Free); TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_Free);