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

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_ShapeSize);
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_HloModule);
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(
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) {
output.AddToBeReleased(
@ -170,7 +171,8 @@ class TpuExecutable : public TpuExecutableInterface {
.Release()
.value());
}
delete[] se_execution_output.to_be_released;
ExecutorApiFn()->TpuExecutable_FreeMaybeOwningDeviceMemoryArrayFn(
se_execution_output.to_be_released);
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_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,
const char** fingerprint,
size_t* size);
@ -468,6 +480,8 @@ struct TfTpu_ExecutorApiFn {
TFTPU_ADD_FN_IN_STRUCT(TpuCompiler_Compile);
TFTPU_ADD_FN_IN_STRUCT(TpuCompiler_ShapeSize);
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_HloModule);
TFTPU_ADD_FN_IN_STRUCT(TpuExecutable_Free);