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

Make StatusHelper safer to use and fix a related memory leak.

Browse Source 
PiperOrigin-RevId: 323124482
Change-Id: I8d3528dd59a1ba1a9f85bf07c93a46eb32f7e645
This commit is contained in:
Wenhao Jia 2020-07-24 21:40:18 -07:00 committed by TensorFlower Gardener
parent 4874f237b8
commit 0de929fdc1
2 changed files with 20 additions and 15 deletions
tensorflow/stream_executor/tpu
status_helper.htpu_transfer_manager.cc

20
tensorflow/stream_executor/tpu/status_helper.h
View File

@ -29,22 +29,24 @@ class StatusHelper {
tensorflow::tpu::ExecutorApiFn()->TpuStatus_FreeFn(c_status);
}
bool ok() const {
return tensorflow::tpu::ExecutorApiFn()->TpuStatus_CodeFn(c_status) == 0;
}
tensorflow::Status status() const {
if (!ok()) {
static tensorflow::Status FromC(SE_Status* const c_status) {
if (tensorflow::tpu::ExecutorApiFn()->TpuStatus_OkFn(c_status)) {
return tensorflow::Status::OK();
} else {
return tensorflow::Status(
tensorflow::error::Code(
tensorflow::tpu::ExecutorApiFn()->TpuStatus_CodeFn(c_status)),
tensorflow::tpu::ExecutorApiFn()->TpuStatus_MessageFn(c_status));
} else {
return tensorflow::Status::OK();
}
}
SE_Status* c_status; // NOLINT
bool ok() const {
return tensorflow::tpu::ExecutorApiFn()->TpuStatus_OkFn(c_status);
}
tensorflow::Status status() const { return FromC(c_status); }
SE_Status* const c_status; // NOLINT
};
#endif // TENSORFLOW_STREAM_EXECUTOR_TPU_STATUS_HELPER_H_

15
tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
View File

@ -15,6 +15,8 @@ limitations under the License.
#include "tensorflow/stream_executor/tpu/tpu_transfer_manager.h"
#include <utility>
#include "tensorflow/compiler/xla/shape_util.h"
#include "tensorflow/compiler/xla/xla_data.pb.h"
#include "tensorflow/core/tpu/tpu_api.h"
@ -80,19 +82,20 @@ Status TpuTransferManager::TransferLiteralToDeviceAsync(
struct TransferFromDeviceState {
std::atomic<int64_t> remaining_transfers;
StatusHelper status_helper;
SE_Status* overall_status =
tpu::ExecutorApiFn()->TpuStatus_NewFn(); // OK or the first error
std::function<void(Status)> done;
void TransferFinished(SE_Status* status) {
if (!tpu::ExecutorApiFn()->TpuStatus_OkFn(status) &&
tpu::ExecutorApiFn()->TpuStatus_OkFn(status_helper.c_status)) {
status_helper.c_status = status;
} else {
tpu::ExecutorApiFn()->TpuStatus_FreeFn(status);
tpu::ExecutorApiFn()->TpuStatus_OkFn(overall_status)) {
std::swap(overall_status, status);
}
tpu::ExecutorApiFn()->TpuStatus_FreeFn(status);
if (--remaining_transfers == 0) {
done(status_helper.status());
done(StatusHelper::FromC(overall_status));
tpu::ExecutorApiFn()->TpuStatus_FreeFn(overall_status);
delete this;
}
}