| 397 | |
| 398 | MGB_DYN_TYPE_OBJ_FINAL_IMPL(TensorRTOpr); |
| 399 | TensorRTOpr::TensorRTOpr( |
| 400 | std::shared_ptr<nvinfer1::IBuilder> builder, |
| 401 | std::shared_ptr<nvinfer1::INetworkDefinition> network, |
| 402 | TensorRTGraphFeatureBits feature_bits, |
| 403 | std::shared_ptr<GpuAllocator> gpu_allocator, const VarNodeArray& inputs, |
| 404 | std::shared_ptr<nvinfer1::ICudaEngine> engine, const OperatorNodeConfig& config) |
| 405 | : Super(inputs.at(0)->owner_graph(), config, "tensor_rt", {inputs.at(0)}), |
| 406 | m_builder{std::move(builder)}, |
| 407 | m_gpu_allocator{std::move(gpu_allocator)}, |
| 408 | m_network{std::move(network)}, |
| 409 | m_engine{std::move(engine)}, |
| 410 | m_feature_bits{feature_bits} { |
| 411 | mgb_assert( |
| 412 | inputs[0]->comp_node().device_type() == CompNode::DeviceType::CUDA, |
| 413 | "TensorRTOpr can only be used on cuda comp nodes; got %s", |
| 414 | inputs[0]->comp_node().to_string().c_str()); |
| 415 | mgb_assert( |
| 416 | inputs.size() == static_cast<size_t>(m_network->getNbInputs()), |
| 417 | "inputs size not equal: expect=%zu got=%d", inputs.size(), |
| 418 | m_network->getNbInputs()); |
| 419 | for (auto i : inputs) { |
| 420 | add_input({i}); |
| 421 | } |
| 422 | if (m_network->getNbOutputs() == 1) |
| 423 | add_output(None); |
| 424 | else { |
| 425 | for (int i = 0; i < m_network->getNbOutputs(); ++i) |
| 426 | add_output(ssprintf("o%d", i)); |
| 427 | } |
| 428 | cg::add_workspace_output(this); |
| 429 | |
| 430 | add_equivalence_component<mgb::ScalarHash<void*>>(m_network.get()); |
| 431 | mgb_assert(m_builder != nullptr); |
| 432 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 433 | m_builder_config = { |
| 434 | m_builder->createBuilderConfig(), |
| 435 | TensorRTDeleter<nvinfer1::IBuilderConfig>()}; |
| 436 | m_builder_config->setMaxWorkspaceSize(1 << 30); |
| 437 | if (m_feature_bits == TensorRTGraphFeatureBits::NCHW4_QINT8) { |
| 438 | mgb_assert( |
| 439 | m_builder->platformHasFastInt8(), |
| 440 | "Cuda platform does not support fast native int8"); |
| 441 | m_builder_config->setInt8Calibrator(nullptr); |
| 442 | nvinfer1::BuilderFlags flags; |
| 443 | flags = 1 << static_cast<int>(nvinfer1::BuilderFlag::kINT8); |
| 444 | m_builder_config->setFlags(flags); |
| 445 | } |
| 446 | #else |
| 447 | m_builder->setMaxWorkspaceSize(1 << 30); |
| 448 | if (m_feature_bits == TensorRTGraphFeatureBits::NCHW4_QINT8) { |
| 449 | // check has fast int8 |
| 450 | m_builder->setInt8Mode(true); |
| 451 | m_builder->setInt8Calibrator(nullptr); |
| 452 | m_builder->setStrictTypeConstraints(false); |
| 453 | } |
| 454 | #endif |
| 455 | if (!m_gpu_allocator) { |
| 456 | m_gpu_allocator = std::make_shared<GpuAllocator>(inputs[0]->comp_node()); |
nothing calls this directly
no test coverage detected