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

Function apply_on_var_node

imperative/src/impl/ops/collective_comm.cpp:17–37  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

15#if MGB_ENABLE_OPR_MM
16namespace {
17cg::OperatorNodeBase* apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
18 auto&& comm = def.cast_final_safe<CollectiveComm>();
19 auto group_client = std::make_shared<opr::GroupClientProxy>(
20 ssprintf("%s:%d", comm.addr.data(), comm.port));
21 SmallVector<std::shared_ptr<mgb::DeviceTensorND>> dev_buffer_arr(1, nullptr);
22 auto disable = std::make_shared<DTypeScalar>();
23 disable->set(0);
24
25 OperatorNodeConfig config{comm.make_name()};
26 if (comm.comp_node.size() > 0) {
27 config.comp_node(CompNode::load(comm.comp_node));
28 }
29
30 mgb_assert(inputs.size() == 1, "exactly one input expected");
31 auto&& graph = inputs[0]->owner_graph();
32
33 return graph->insert_opr(std::make_unique<opr::CollectiveComm>(
34 inputs, graph, comm.key, comm.nr_devices, comm.is_root, comm.rank,
35 comm.local_grad, group_client, comm.mode, comm.dtype, comm.backend,
36 dev_buffer_arr, config, disable));
37}
38
39std::tuple<std::string, std::string> split_address(
40 const std::string& address_and_port) {

Callers

nothing calls this directly

Calls 8

loadFunction · 0.50
dataMethod · 0.45
setMethod · 0.45
make_nameMethod · 0.45
sizeMethod · 0.45
comp_nodeMethod · 0.45
owner_graphMethod · 0.45
insert_oprMethod · 0.45

Tested by

no test coverage detected