| 842 | } |
| 843 | |
| 844 | uint64_t PackAllReduceReplacePass::collect_groups( |
| 845 | OperatorNodeBase* opr, |
| 846 | ThinHashMap<uint64_t, std::shared_ptr<GroupInfo>>& group_info, |
| 847 | ThinHashMap<uint64_t, cg::OprNodeArray>& groups) { |
| 848 | // check CollectiveComm oprs that have been marked in PackAllReduceScanPass |
| 849 | if (!opr->same_type<opr::CollectiveComm>()) |
| 850 | return 0; |
| 851 | opr::CollectiveComm& comm = opr->cast_final_safe<opr::CollectiveComm>(); |
| 852 | if (comm.pack_hash() == 0) |
| 853 | return 0; // pack_hash not set |
| 854 | |
| 855 | VarNode* var = comm.input(0); |
| 856 | auto info = std::make_shared<GroupInfo>( |
| 857 | var->comp_node().locator().device, var->dtype(), comm.nr_devices(), |
| 858 | comm.is_root(), comm.rank(), comm.group_client(), comm.backend()); |
| 859 | uint64_t hash = info->hash(comm.pack_hash()); |
| 860 | if (group_info.find(hash) == group_info.end()) { |
| 861 | group_info.emplace(hash, info); |
| 862 | } |
| 863 | groups[hash].push_back(opr); |
| 864 | return hash; |
| 865 | } |
| 866 | |
| 867 | void PackAllReduceReplacePass::divide_packs( |
| 868 | const ThinHashMap<uint64_t, cg::OprNodeArray>& groups, |
nothing calls this directly
no test coverage detected