| 669 | } |
| 670 | |
| 671 | void TensorRTOpr::build_engine_from_cache() { |
| 672 | TensorRTUniquePtr<nvinfer1::IRuntime> runtime{ |
| 673 | nvinfer1::createInferRuntime(TensorRTOpr::Logger::instance()), {}}; |
| 674 | runtime->setGpuAllocator(m_gpu_allocator.get()); |
| 675 | auto ret = TensorRTEngineCache::inst().get( |
| 676 | TensorRTEngineCache::make_key_from_trt_opr(this)); |
| 677 | if (!ret.valid()) |
| 678 | return; |
| 679 | comp_node().activate(); |
| 680 | #if NV_TENSOR_RT_VERSION < 8000 |
| 681 | auto engine = runtime->deserializeCudaEngine( |
| 682 | reinterpret_cast<const void*>(ret->ptr), ret->size, nullptr); |
| 683 | #else |
| 684 | auto engine = runtime->deserializeCudaEngine( |
| 685 | reinterpret_cast<const void*>(ret->ptr), ret->size); |
| 686 | #endif |
| 687 | mgb_assert(engine, "failed to deserialize ICudaEngine"); |
| 688 | m_engine = {engine, TensorRTDeleter<nvinfer1::ICudaEngine>()}; |
| 689 | } |
| 690 | |
| 691 | void TensorRTOpr::serialize_engine_to_cache() const { |
| 692 | TensorRTUniquePtr<nvinfer1::IHostMemory> buf{trt_cuda_engine()->serialize(), {}}; |