| 18 | using namespace tensorrt; |
| 19 | |
| 20 | void TrtReplaceChecker::ensure_init_graph() { |
| 21 | if (m_trt_y.node()) |
| 22 | return; |
| 23 | |
| 24 | SymbolVarArray inputs(m_nr_input); |
| 25 | for (size_t i = 0; i < m_nr_input; ++i) { |
| 26 | if (m_mark_inp_const.count(i)) { |
| 27 | inputs[i] = opr::SharedDeviceTensor::make(*m_graph, *m_inputs_val[i]) |
| 28 | .rename(ssprintf("inp%zu", i)); |
| 29 | } else { |
| 30 | inputs[i] = opr::Host2DeviceCopy::make(*m_graph, m_inputs_val[i]) |
| 31 | .rename(ssprintf("inp%zu", i)); |
| 32 | } |
| 33 | |
| 34 | auto dt = m_idx2dtype.find(i); |
| 35 | if (dt != m_idx2dtype.end()) { |
| 36 | inputs[i] = opr::TypeCvt::make(inputs[i], dt->second); |
| 37 | } |
| 38 | } |
| 39 | m_truth_y = m_exp_func(inputs); |
| 40 | |
| 41 | ComputingGraph::Options opt; |
| 42 | opt.graph_opt_level = 0; |
| 43 | unpack_vector( |
| 44 | gopt::GraphOptimizer{} |
| 45 | .add_pass<gopt::ExpandFusedArithPass>() |
| 46 | .add_pass<gopt::TensorRTReplacePass>() |
| 47 | .add_pass<gopt::ArithFusePass>() |
| 48 | .apply({{m_truth_y}}) |
| 49 | .endpoint_vars(), |
| 50 | m_trt_y); |
| 51 | |
| 52 | size_t nr_trt_opr = 0; |
| 53 | cg::DepOprIter{[&nr_trt_opr, this](cg::OperatorNodeBase* opr) { |
| 54 | if (opr->same_type<TensorRTOpr>()) { |
| 55 | ++nr_trt_opr; |
| 56 | } |
| 57 | }}.add(m_trt_y.node()); |
| 58 | mgb_assert(nr_trt_opr >= 1); |
| 59 | |
| 60 | ComputingGraph::OutputSpec outspec(2); |
| 61 | outspec[0] = make_callback_copy(m_truth_y, std::get<0>(m_output_val), false); |
| 62 | outspec[1] = make_callback_copy(m_trt_y, std::get<1>(m_output_val), false); |
| 63 | |
| 64 | m_graph->options().graph_opt.tensorrt = false; |
| 65 | m_func = m_graph->compile(outspec); |
| 66 | } |
| 67 | |
| 68 | TrtReplaceChecker& TrtReplaceChecker::run(const TensorShapeArray& input_shapes) { |
| 69 | if (::testing::Test::HasFailure()) { |