| 865 | } |
| 866 | |
| 867 | void PackAllReduceReplacePass::divide_packs( |
| 868 | const ThinHashMap<uint64_t, cg::OprNodeArray>& groups, |
| 869 | ThinHashMap<uint64_t, std::vector<cg::OprNodeArray>>& packs, size_t max_size) { |
| 870 | cg::OprNodeArray pack; |
| 871 | size_t sum = 0; |
| 872 | for (auto it : groups) { |
| 873 | uint64_t hash = it.first; |
| 874 | const cg::OprNodeArray& group = it.second; |
| 875 | for (size_t i = 0; i < group.size(); i++) { |
| 876 | OperatorNodeBase* opr = group[i]; |
| 877 | VarNode* var = opr->input(0); |
| 878 | const TensorShape* shape = |
| 879 | var->owner_graph()->static_infer_manager().infer_shape_fallible( |
| 880 | var); |
| 881 | if (shape == nullptr) |
| 882 | continue; |
| 883 | pack.push_back(opr); |
| 884 | sum += var->dtype().size(shape->total_nr_elems()); |
| 885 | if (sum >= max_size) { |
| 886 | if (pack.size() > 1) |
| 887 | packs[hash].push_back(pack); |
| 888 | pack.clear(); |
| 889 | sum = 0; |
| 890 | } |
| 891 | } |
| 892 | if (pack.size() > 1) |
| 893 | packs[hash].push_back(pack); |
| 894 | pack.clear(); |
| 895 | sum = 0; |
| 896 | } |
| 897 | } |
| 898 | |
| 899 | void PackAllReduceReplacePass::insert_packed_oprs( |
| 900 | size_t pack_id, const cg::OprNodeArray& pack, std::shared_ptr<GroupInfo> info, |
nothing calls this directly
no test coverage detected