| 289 | #endif |
| 290 | |
| 291 | void TensorRTManager::exec( |
| 292 | cg::SingleCNOperatorNodeBase* opr, CompNode comp_node_check, |
| 293 | nvinfer1::ICudaEngine* engine, size_t batch, bool use_trt_profiler) { |
| 294 | auto comp_node = opr->comp_node(); |
| 295 | // ICudaEngine is bound to the currently active device |
| 296 | comp_node.activate(); |
| 297 | |
| 298 | if (comp_node_check.valid()) { |
| 299 | mgb_assert( |
| 300 | comp_node_check == comp_node, |
| 301 | "gpu allocator is on %s, but execution is on %s", |
| 302 | comp_node_check.to_string().c_str(), comp_node.to_string().c_str()); |
| 303 | } |
| 304 | auto workspace_ptr = opr->output().back()->dev_tensor().raw_ptr(); |
| 305 | bool should_reinit_device_memory = |
| 306 | !m_context || (m_device_workspace_memory_ptr != workspace_ptr) || |
| 307 | (workspace_ptr == nullptr); |
| 308 | if (!m_context) { |
| 309 | TensorShapeArray arr; |
| 310 | for (auto&& i : opr->input()) { |
| 311 | arr.push_back(i->shape()); |
| 312 | } |
| 313 | create_trt_context(comp_node, arr, engine); |
| 314 | } |
| 315 | m_trt_iobuf.resize(engine->getNbBindings()); |
| 316 | bool is_trt_opr = false; |
| 317 | if (opr->same_type<TensorRTOpr>()) { |
| 318 | is_trt_opr = true; |
| 319 | auto network = opr->cast_final_safe<TensorRTOpr>().trt_network_def(); |
| 320 | int nr_input = network->getNbInputs(); |
| 321 | for (int i = 0; i < nr_input; ++i) { |
| 322 | int binding_idx = engine->getBindingIndex(network->getInput(i)->getName()); |
| 323 | m_trt_iobuf[binding_idx] = opr->input(i)->dev_tensor().raw_ptr(); |
| 324 | } |
| 325 | int nr_output = network->getNbOutputs(); |
| 326 | for (int i = 0; i < nr_output; ++i) { |
| 327 | int binding_idx = engine->getBindingIndex(network->getOutput(i)->getName()); |
| 328 | m_trt_iobuf[binding_idx] = opr->output(i)->dev_tensor().raw_ptr(); |
| 329 | } |
| 330 | } else { |
| 331 | for (size_t i = 0; i < opr->input().size(); ++i) { |
| 332 | m_trt_iobuf[i + m_offset] = opr->input(i)->dev_tensor().raw_ptr(); |
| 333 | } |
| 334 | for (size_t i = 0; i < opr->output().size() - 1; ++i) { |
| 335 | m_trt_iobuf[opr->input().size() + i + m_offset] = |
| 336 | opr->output(i)->dev_tensor().raw_ptr(); |
| 337 | } |
| 338 | } |
| 339 | MGB_MARK_USED_VAR(is_trt_opr); |
| 340 | #if TENSOR_RT_MANAGE_ALL_WORKSPACE |
| 341 | MGB_MARK_USED_VAR(should_reinit_device_memory); |
| 342 | #else |
| 343 | if (should_reinit_device_memory) { |
| 344 | mgb_assert( |
| 345 | opr->output().back()->shape()[0] == intl::workspace_size(engine) && |
| 346 | !(reinterpret_cast<uintptr_t>(workspace_ptr) % 256)); |
| 347 | m_context->setDeviceMemory(workspace_ptr); |
| 348 | m_device_workspace_memory_ptr = workspace_ptr; |
no test coverage detected