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