Skip to content

Commit

Permalink
fix: Correct mutex scope in execute_engine() (#3310)
Browse files Browse the repository at this point in the history
  • Loading branch information
keehyuna authored Dec 16, 2024
1 parent b0b6f38 commit bed5d37
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions core/runtime/execute_engine.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
LOG_DEBUG(
"Attempting to run engine (ID: " << compiled_engine->name
<< "); Hardware Compatible: " << compiled_engine->hardware_compatible);

// nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex for it.
// Other IExecutionContext methods and runtime states should be in same scope as well
std::unique_lock<std::mutex> lock(compiled_engine->mu);
if (compiled_engine->profile_execution) {
std::stringstream ss;
ss << "Execution profiling is enabled, find results here:" << std::endl;
Expand Down Expand Up @@ -307,9 +309,6 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
compiled_engine->engine_stream = c10::cuda::getStreamFromPool(false, current_device_id);
}

// nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex for it.
std::unique_lock<std::mutex> lock(compiled_engine->mu);

{ // Engine Execution (execute on engine stream)
c10::cuda::CUDAStreamGuard stream_guard(compiled_engine->engine_stream);

Expand Down

0 comments on commit bed5d37

Please sign in to comment.