MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / init_output_comp_node

Method init_output_comp_node

src/opr/impl/tensor_manip.cpp:1043–1082  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1041}
1042
1043void Split::init_output_comp_node() {
1044 auto&& conf_node = config().comp_node();
1045 auto&& cn_opt = owner_graph()->seq_comp_node_optimizer();
1046
1047 // details of each comp_node specified
1048 if (conf_node.size() > 1) {
1049 mgb_assert(
1050 conf_node.size() == output().size(),
1051 "number of CompNodes specified in config should equal to number"
1052 " of output, but got %zd configured CompNodes while there are"
1053 " %zd output (node_name=%s node_type=%s)",
1054 conf_node.size(), output().size(), cname(), dyn_typeinfo()->name);
1055 auto cn0 = input(0)->comp_node();
1056 for (size_t i = 0; i < output().size(); i++) {
1057 auto dvar = output(i);
1058 dvar->comp_node(conf_node[i]);
1059 if (conf_node[i].mem_node() != cn0.mem_node())
1060 cn_opt.register_stream_var(
1061 dvar, {CompNode::Stream::COPY,
1062 cg::SeqCompNodeOptimizer::StreamPropType::WEAK});
1063 }
1064 return;
1065 }
1066
1067 CompNode cn;
1068 if (conf_node.size() == 1) {
1069 cn = conf_node[0];
1070 } else {
1071 cn = input(0)->comp_node();
1072 }
1073 for (auto i : output())
1074 i->comp_node(cn);
1075
1076 if (cn.mem_node() != input(0)->comp_node().mem_node()) {
1077 for (auto i : output())
1078 cn_opt.register_stream_var(
1079 i, {CompNode::Stream::COPY,
1080 cg::SeqCompNodeOptimizer::StreamPropType::WEAK});
1081 }
1082}
1083
1084cg::OperatorNodeBase::NodeProp* Split::do_make_node_prop() const {
1085 auto rst = OperatorNodeBase::do_make_node_prop();

Callers

nothing calls this directly

Calls 5

init_output_comp_nodeFunction · 0.50
comp_nodeMethod · 0.45
sizeMethod · 0.45
mem_nodeMethod · 0.45
register_stream_varMethod · 0.45

Tested by

no test coverage detected