| 88 | } // namespace |
| 89 | |
| 90 | bool Optimizer::FusionAndShard( |
| 91 | std::shared_ptr<FusionAndShardRule> rule, DAGDef* dag) { |
| 92 | std::string target_name = rule->target_name_; |
| 93 | bool dynamic_split = rule->dynamic_split_; |
| 94 | bool dynamic_output = rule->dynamic_output_; |
| 95 | bool fusion_success = true; |
| 96 | bool one_shot = true; |
| 97 | while (fusion_success) { |
| 98 | /* find pattern */ |
| 99 | std::vector<std::unordered_map<int32_t, int32_t>> patterns; |
| 100 | if (rule->opt_type_ == graph_partition && one_shot) { |
| 101 | // build pattern directly |
| 102 | std::unordered_map<int32_t, std::shared_ptr<NodeDef>> sub_dag = |
| 103 | rule->sub_dag_.GetNodeMap(); |
| 104 | std::unordered_map<int32_t, int32_t> pattern; |
| 105 | for (auto it = sub_dag.begin(); it != sub_dag.end(); ++it) { |
| 106 | pattern[it->first] = it->first; |
| 107 | } |
| 108 | patterns.push_back(pattern); |
| 109 | one_shot = false; |
| 110 | } else if (rule->opt_type_ != graph_partition) { |
| 111 | patterns = SubGraphMatch(*dag, rule->sub_dag_, rule->extra_cond_); |
| 112 | } |
| 113 | |
| 114 | fusion_success = false; |
| 115 | /* try to fusion and shard */ |
| 116 | for (const std::unordered_map<int32_t, int32_t>& pattern : patterns) { |
| 117 | std::unordered_set<int32_t> fusion_set; |
| 118 | PrepareFusionShardRule( |
| 119 | dynamic_split, dynamic_output, *dag, pattern, rule, &fusion_set); |
| 120 | FusionRule fusion_rule( |
| 121 | target_name, pattern, rule->fusion_output_map_, rule->fusion_nodes_); |
| 122 | // fusion |
| 123 | int32_t fusion_node_id = dag->FusionNodes( |
| 124 | fusion_set, fusion_rule); |
| 125 | fusion_success = fusion_success || fusion_node_id != -1; |
| 126 | // shard |
| 127 | if (fusion_success) { |
| 128 | if (target_name == "REMOTE") { // need to be shard |
| 129 | ShardRule shard_rule(rule->split_op_info_, |
| 130 | rule->merge_op_info_, |
| 131 | rule->split_num_); |
| 132 | if (!dag->ShardRemoteNodeDef(fusion_node_id, shard_rule)) { |
| 133 | return false; |
| 134 | } |
| 135 | } |
| 136 | break; |
| 137 | } |
| 138 | } |
| 139 | } |
| 140 | return true; |
| 141 | } |
| 142 | |
| 143 | bool Optimizer::UniqueAndGather( |
| 144 | std::shared_ptr<UniqueAndGatherRule> rule, DAGDef* dag) { |
nothing calls this directly
no test coverage detected