TFLite Metal delegate: Memory leak when aggressive wait mode is used has been fixed.

PiperOrigin-RevId: 308749950
Change-Id: I5786e1673984df9c22399726b2b638670f9b1d5c
This commit is contained in:
A. Unique TensorFlower 2020-04-27 20:10:17 -07:00 committed by TensorFlower Gardener
parent b620b4dfcb
commit 2a72ad4071

View File

@ -120,19 +120,21 @@ class GpuAlarmClock {
alarm_thread_ = std::thread([this]() {
id<MTLCommandBuffer> prev_command_buffer;
while (!release_thread_) {
if (active_alarms_ == total_alarms_) {
id<MTLCommandBuffer> command_buffer = [command_queue_ commandBuffer];
id<MTLComputeCommandEncoder> encoder = [command_buffer computeCommandEncoder];
[encoder setComputePipelineState:stub_program_];
[encoder setBuffer:stub_buffer_ offset:0 atIndex:0];
[encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
[encoder endEncoding];
[command_buffer commit];
if (prev_command_buffer != nil) [prev_command_buffer waitUntilScheduled];
prev_command_buffer = command_buffer;
} else {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
@autoreleasepool {
if (active_alarms_ == total_alarms_) {
id<MTLCommandBuffer> command_buffer = [command_queue_ commandBuffer];
id<MTLComputeCommandEncoder> encoder = [command_buffer computeCommandEncoder];
[encoder setComputePipelineState:stub_program_];
[encoder setBuffer:stub_buffer_ offset:0 atIndex:0];
[encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
[encoder endEncoding];
[command_buffer commit];
if (prev_command_buffer != nil) [prev_command_buffer waitUntilScheduled];
prev_command_buffer = command_buffer;
} else {
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
}
}
});