Skip to content

Commit bed5d37

Browse files
authored
fix: Correct mutex scope in execute_engine() (#3310)
1 parent b0b6f38 commit bed5d37

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

core/runtime/execute_engine.cpp

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
9898
LOG_DEBUG(
9999
"Attempting to run engine (ID: " << compiled_engine->name
100100
<< "); Hardware Compatible: " << compiled_engine->hardware_compatible);
101-
101+
// nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex for it.
102+
// Other IExecutionContext methods and runtime states should be in same scope as well
103+
std::unique_lock<std::mutex> lock(compiled_engine->mu);
102104
if (compiled_engine->profile_execution) {
103105
std::stringstream ss;
104106
ss << "Execution profiling is enabled, find results here:" << std::endl;
@@ -307,9 +309,6 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
307309
compiled_engine->engine_stream = c10::cuda::getStreamFromPool(false, current_device_id);
308310
}
309311

310-
// nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex for it.
311-
std::unique_lock<std::mutex> lock(compiled_engine->mu);
312-
313312
{ // Engine Execution (execute on engine stream)
314313
c10::cuda::CUDAStreamGuard stream_guard(compiled_engine->engine_stream);
315314

0 commit comments

Comments
 (0)