| 46 | |
| 47 | MGB_DYN_TYPE_OBJ_FINAL_IMPL(TensorRTRuntimeOpr); |
| 48 | TensorRTRuntimeOpr::TensorRTRuntimeOpr( |
| 49 | std::shared_ptr<nvinfer1::ICudaEngine> engine, |
| 50 | std::shared_ptr<GpuAllocator> gpu_allocator, const VarNodeArray& inputs, |
| 51 | const OperatorNodeConfig& config) |
| 52 | : Super(inputs.at(0)->owner_graph(), config, "tensor_rt", {inputs.at(0)}), |
| 53 | m_gpu_allocator{std::move(gpu_allocator)}, |
| 54 | m_engine{std::move(engine)}, |
| 55 | m_trt_engine_has_batch{false} { |
| 56 | mgb_assert( |
| 57 | inputs[0]->comp_node().device_type() == CompNode::DeviceType::CUDA, |
| 58 | "TensorRTRuntimeOpr can only be used on cuda comp nodes; got %s", |
| 59 | inputs[0]->comp_node().to_string().c_str()); |
| 60 | size_t nr_input = 0; |
| 61 | bool is_input = true; |
| 62 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 63 | auto profile_num = m_engine->getNbOptimizationProfiles(); |
| 64 | #else |
| 65 | int profile_num = 1; |
| 66 | #endif |
| 67 | auto bindings_per_profile = m_engine->getNbBindings() / profile_num; |
| 68 | for (int i = 0; i < bindings_per_profile; ++i) { |
| 69 | if (m_engine->bindingIsInput(nr_input)) { |
| 70 | mgb_assert(is_input, "mixed input/output bindings"); |
| 71 | // nbDims == 3, means CHW, without batch |
| 72 | if (m_engine->getBindingDimensions(nr_input).nbDims != 3) |
| 73 | m_trt_engine_has_batch = true; |
| 74 | ++nr_input; |
| 75 | } else { |
| 76 | is_input = false; |
| 77 | } |
| 78 | } |
| 79 | size_t nr_output = bindings_per_profile - nr_input; |
| 80 | mgb_assert( |
| 81 | nr_input == inputs.size(), "inputs size not equal: expect=%zu got=%zu", |
| 82 | nr_input, inputs.size()); |
| 83 | for (auto i : inputs) { |
| 84 | add_input({i}); |
| 85 | } |
| 86 | if (nr_output == 1) { |
| 87 | add_output(None); |
| 88 | } else { |
| 89 | for (size_t i = 0; i < nr_output; ++i) |
| 90 | add_output(ssprintf("o%zu", i)); |
| 91 | } |
| 92 | cg::add_workspace_output(this); |
| 93 | add_equivalence_component<mgb::ScalarHash<void*>>(m_engine.get()); |
| 94 | } |
| 95 | |
| 96 | void TensorRTRuntimeOpr::get_output_var_shape( |
| 97 | const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const { |
nothing calls this directly
no test coverage detected