| 710 | } |
| 711 | |
| 712 | void LoopGrad::GradDesc::on_sub_graph_func_compile( |
| 713 | ComputingGraph::OutputSpec& out_spec) { |
| 714 | { |
| 715 | // append extra targets to out_spec |
| 716 | size_t idx = 0, nr_out_spec = out_spec.size(); |
| 717 | mgb_assert( |
| 718 | out_spec[idx++].first.node() == |
| 719 | loop_cond_manager().subgraph_outspec_item().first.node()); |
| 720 | |
| 721 | for (auto&& i : output_record_spec()) { |
| 722 | if (!i.enabled()) |
| 723 | continue; |
| 724 | auto&& spec = out_spec[idx++]; |
| 725 | mgb_assert(spec.first.node() == i.var_sub()); |
| 726 | i.recorder() |
| 727 | ->cast_final_safe<OutputRecorderSumIntoDest>() |
| 728 | .add_extra_compile_output_spec(out_spec); |
| 729 | } |
| 730 | mgb_assert(idx == nr_out_spec); |
| 731 | |
| 732 | // add outspec for AssignorGradOpr |
| 733 | auto cb = [&](OperatorNodeBase* opr) { |
| 734 | if (opr->same_type<AssignorGradOpr>()) { |
| 735 | opr->cast_final<AssignorGradOpr>().add_extra_compile_output_spec( |
| 736 | out_spec); |
| 737 | } |
| 738 | }; |
| 739 | cg::DepOprIter iter{cb}; |
| 740 | for (idx = 0, nr_out_spec = out_spec.size(); idx < nr_out_spec; ++idx) { |
| 741 | iter.add(out_spec[idx].first.node()->owner_opr()); |
| 742 | } |
| 743 | } |
| 744 | int opt_level = owner_graph()->options().graph_opt_level; |
| 745 | if (std::abs(opt_level) < 2) |
| 746 | return; |
| 747 | VarNodeArray endpoints; |
| 748 | endpoints.reserve(out_spec.size()); |
| 749 | endpoints.push_back(m_orig_loop_cond_var.node()); |
| 750 | for (size_t i = 1; i < out_spec.size(); ++i) |
| 751 | endpoints.push_back(out_spec[i].first.node()); |
| 752 | |
| 753 | if (endpoints == m_prev_sub_graph_opt_endpoints_inp) { |
| 754 | endpoints = m_prev_sub_graph_opt_endpoints_out; |
| 755 | } else { |
| 756 | m_prev_sub_graph_opt_endpoints_inp = endpoints; |
| 757 | gopt::GraphOptimizer() |
| 758 | .verbosity(0) |
| 759 | .add_preset_passes() |
| 760 | .enable_check_result(opt_level < 0) |
| 761 | .apply_inplace(endpoints); |
| 762 | m_prev_sub_graph_opt_endpoints_out = endpoints; |
| 763 | |
| 764 | // NO_MEM_RECLAIM flag is required for OutputRecorderSumIntoDest |
| 765 | for (size_t i = 0; i < endpoints.size(); ++i) { |
| 766 | constexpr auto F = VarNode::Flag::NO_MEM_RECLAIM; |
| 767 | if (m_prev_sub_graph_opt_endpoints_inp[i]->contain_flag(F)) { |
| 768 | endpoints[i]->add_flag(F); |
| 769 | } |
nothing calls this directly
no test coverage detected