| 635 | } |
| 636 | |
| 637 | void CollectiveComm::do_execute(ExecEnv& env) { |
| 638 | auto&& trait = ModeTrait::from_mode(m_param.mode); |
| 639 | mgb_assert( |
| 640 | owner_graph()->options().async_exec_level, |
| 641 | "collective comm must be used with async dispatch"); |
| 642 | mgb_assert(output().size() == 1, "collective comm only support exactly one output"); |
| 643 | |
| 644 | auto disable = m_disable->get_cast<int>(); |
| 645 | if (disable == 1) |
| 646 | return; |
| 647 | mgb_assert( |
| 648 | disable == 0, |
| 649 | "disable flag on CollectiveComm can only be 0 or 1," |
| 650 | " got %d actually.", |
| 651 | disable); |
| 652 | |
| 653 | auto cn = output(0)->comp_node(); |
| 654 | auto runner = [this, cn, &trait] { |
| 655 | opr_register(); |
| 656 | cn.activate(); |
| 657 | |
| 658 | if (m_debug_mode) { |
| 659 | mgb_log_debug( |
| 660 | "collective comm: executing %s, rank = %d, key = %s", cname(), |
| 661 | rank(), key().c_str()); |
| 662 | } |
| 663 | |
| 664 | owner_graph()->event().signal_inplace<cg::event::BeforeKernel>(this, cn); |
| 665 | trait.exec(this); |
| 666 | owner_graph()->event().signal_inplace<cg::event::AfterKernel>(this, cn); |
| 667 | }; |
| 668 | env.dispatch_on_comp_node(cn, runner); |
| 669 | } |
| 670 | |
| 671 | void CollectiveComm::on_output_comp_node_stream_changed() {} |
| 672 | |