| 15 | #if MGB_ENABLE_OPR_MM |
| 16 | namespace { |
| 17 | cg::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 | |
| 39 | std::tuple<std::string, std::string> split_address( |
| 40 | const std::string& address_and_port) { |
nothing calls this directly
no test coverage detected