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

Function make_from_computing_graph

imperative/src/impl/subgraph_detail.cpp:176–245  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

174}
175
176EncodedSubgraph make_from_computing_graph(
177 const VarNodeArray& inputs, const VarNodeArray& outputs) {
178 Subgraph subgraph;
179 std::unordered_map<VarNode*, size_t> var2idx;
180 size_t next_idx = 0;
181 var2idx[nullptr] = next_idx++;
182 for (auto&& input : inputs) {
183 if (input) {
184 var2idx[input] = next_idx++;
185 }
186 }
187 auto is_tensor_holder = [](cg::OperatorNodeBase* op) {
188 return op->input().empty();
189 };
190 auto as_tensor = [](VarNode* var) -> TensorPtr {
191 auto* opr = var->owner_opr();
192 if (auto* imm_tensor = opr->try_cast_final<opr::ImmutableTensor>()) {
193 auto&& dv = imm_tensor->value();
194 HostTensorND hv(dv.comp_node(), dv.shape(), dv.dtype());
195 // get host value
196 auto&& cpu_value = imm_tensor->host_value();
197 mgb_assert(cpu_value.comp_node() == CompNode::default_cpu());
198 // default_cpu is synchronous with respect to caller
199 hv.proxy_to_default_cpu().copy_from_fixlayout(cpu_value);
200 return Tensor::make(dv, hv);
201 } else if (
202 auto* shared_tensor = opr->try_cast_final<opr::SharedDeviceTensor>()) {
203 return Tensor::make(shared_tensor->get_dev_tensor());
204 } else {
205 mgb_assert(
206 false, "unsupported tensor holder opr %s",
207 opr->dyn_typeinfo()->name);
208 }
209 };
210 cg::DepOprIter iter{[&](cg::OperatorNodeBase* op) {
211 // TODO: implement make_backward_graph for mm ops
212 // mgb_assert(!op->node_prop().contain(cg::OperatorNodeProp::Flag::IMPURE_FUNC));
213 if (is_tensor_holder(op)) {
214 for (auto&& output : op->usable_output()) {
215 subgraph.constants.push_back(
216 {var2idx[output] = next_idx++, as_tensor(output)});
217 }
218 } else {
219 Subgraph::vars_t inputs;
220 Subgraph::vars_t outputs;
221 for (auto&& input : op->input()) {
222 inputs.push_back(var2idx.at(input));
223 }
224 // NOTE: use usable_output
225 for (auto&& output : op->usable_output()) {
226 outputs.push_back(var2idx[output] = next_idx++);
227 }
228 auto opdef = OpDef::make_from_op_node(op);
229 subgraph.exprs.push_back({opdef, inputs, outputs});
230 }
231 }};
232 for (auto&& input : inputs) {
233 if (input) {

Callers 1

make_backward_graphMethod · 0.85

Calls 15

copy_from_fixlayoutMethod · 0.80
get_dev_tensorMethod · 0.80
makeFunction · 0.70
as_tensorFunction · 0.50
make_from_op_nodeFunction · 0.50
emptyMethod · 0.45
inputMethod · 0.45
owner_oprMethod · 0.45
valueMethod · 0.45
comp_nodeMethod · 0.45
shapeMethod · 0.45
dtypeMethod · 0.45

Tested by

no test coverage detected