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

Method shallow_copy

src/opr/impl/loop/forward_sereg.cpp:203–278  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

201}
202
203cg::OperatorNodeBase* LoopSerializer::shallow_copy(
204 const OprShallowCopyContext& orig_ctx, const Loop& opr,
205 const VarNodeArray& inputs, const OperatorNodeConfig& config) {
206 auto orig_desc = static_cast<LoopImpl::FwdDesc*>(opr.m_desc.get());
207 ThinHashMap<VarNode*, size_t> ogvar2inpidx;
208
209 mgb_assert(inputs.size() == opr.input().size());
210 for (size_t i = 0; i < inputs.size(); ++i)
211 ogvar2inpidx[opr.input(i)] = i;
212
213 VarNodeArray cur_opr_inputs;
214 auto varmap_buf = std::make_shared<ThinHashMap<VarNode*, VarNode*>>();
215 auto desc_maker = [&](Loop::Desc& desc) {
216 ThinHashMap<VarNode*, LoopImpl::InputMaker*> assignee2orig_im;
217 auto&& varmap = *varmap_buf;
218
219 // add inputs
220 OprShallowCopyContext ctx{orig_ctx};
221 for (auto inp : orig_desc->all_inputs()) {
222 auto ogvar = inputs.at(ogvar2inpidx.at(inp->orig_var()));
223 auto subvar = desc.add_input(ogvar, inp->param().has_assign);
224 varmap[inp->output(0)] = subvar.node();
225 if (inp->param().has_assign) {
226 assignee2orig_im[subvar.node()] = inp;
227 }
228 ctx.owner_graph(subvar.node()->owner_graph());
229 }
230
231 // copy oprs
232 for (auto opr : orig_desc->sub_graph_oprs()) {
233 if (opr->same_type<LoopImpl::InputMaker>()) {
234 continue;
235 }
236
237 if (opr->same_type<LoopImpl::DescImplBase::CounterProvider>()) {
238 varmap[opr->output(0)] = desc.get_counter_var().node();
239 } else {
240 cur_opr_inputs.clear();
241 for (auto i : opr->input())
242 cur_opr_inputs.push_back(varmap.at(i));
243 auto new_opr =
244 copy_opr_shallow(*opr, cur_opr_inputs, opr->config(), ctx);
245 mgb_assert(new_opr->output().size() == opr->output().size());
246 for (size_t i = 0; i < new_opr->output().size(); ++i)
247 varmap[opr->output(i)] = new_opr->output(i);
248 }
249 }
250 // add outputs in original order
251 for (auto&& i : orig_desc->output_record_spec_no_dedup()) {
252 desc.add_output(varmap.at(i->var_sub()), i->output_mode());
253 }
254 // add assignments
255 for (auto&& i : assignee2orig_im) {
256 desc.assign(i.first, varmap.at(i.second->assignor()));
257 }
258 desc.set_loop_condition(varmap.at(orig_desc->loop_cond_manager().var().node()));
259 };
260

Callers

nothing calls this directly

Calls 15

get_counter_varMethod · 0.80
var_subMethod · 0.80
assignMethod · 0.80
sub_graphMethod · 0.80
makeFunction · 0.70
getMethod · 0.45
sizeMethod · 0.45
inputMethod · 0.45
atMethod · 0.45
add_inputMethod · 0.45
paramMethod · 0.45
outputMethod · 0.45

Tested by

no test coverage detected