| 236 | } |
| 237 | |
| 238 | void AutoParallel::BuildGraph(GraphDef* graph) { |
| 239 | AddSharedNodes(graph); |
| 240 | for (int i = 0; i < num_replicas_; i++) { |
| 241 | AddOneReplica(graph, i); |
| 242 | } |
| 243 | std::set<string> fetches; |
| 244 | for (size_t i = 0; i < item_->fetch.size(); i++) { |
| 245 | for (int j = 0; j < num_replicas_; j++) { |
| 246 | string prefix = strings::StrCat(kAutoParallelPrefix, "-Replica-", j); |
| 247 | string fetch = AddPrefixToNodeName(item_->fetch[i], prefix); |
| 248 | fetches.insert(fetch); |
| 249 | } |
| 250 | } |
| 251 | string name_control = |
| 252 | strings::StrCat(kAutoParallelPrefix, "-Control-", "Fetch"); |
| 253 | auto control = AddNodeControl(name_control, fetches, graph); |
| 254 | |
| 255 | for (const auto& fetch : item_->fetch) { |
| 256 | AddNodeControl(fetch, {control->name()}, graph); |
| 257 | } |
| 258 | *graph->mutable_library() = item_->graph.library(); |
| 259 | *graph->mutable_versions() = item_->graph.versions(); |
| 260 | LOG(INFO) << "Parallelized graph size: " << graph->node_size(); |
| 261 | } |
| 262 | |
| 263 | Status AutoParallel::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 264 | GraphDef* output) { |