| 4199 | } |
| 4200 | |
| 4201 | Status MklLayoutRewritePass::Run(const GraphOptimizationPassOptions& options) { |
| 4202 | if (options.graph == nullptr && options.partition_graphs == nullptr) { |
| 4203 | return Status::OK(); |
| 4204 | } |
| 4205 | if (DisableMKL()) { |
| 4206 | VLOG(2) << "TF-MKL: Disabling MKL"; |
| 4207 | return Status::OK(); |
| 4208 | } |
| 4209 | |
| 4210 | auto process_graph = [&](std::unique_ptr<Graph>* g) { |
| 4211 | // Get the ownership of a graph |
| 4212 | std::unique_ptr<Graph>* ng = std::move(g); |
| 4213 | RunPass(ng); |
| 4214 | // Return the ownership of a graph back |
| 4215 | g->reset(ng->release()); |
| 4216 | }; |
| 4217 | |
| 4218 | if (kMklLayoutRewritePassGroup != |
| 4219 | OptimizationPassRegistry::POST_PARTITIONING) { |
| 4220 | // For any pre-partitioning phase, a graph is stored in options.graph. |
| 4221 | process_graph(options.graph); |
| 4222 | } else { |
| 4223 | // For post partitioning phase, graphs are stored in |
| 4224 | // options.partition_graphs. |
| 4225 | for (auto& pg : *options.partition_graphs) { |
| 4226 | process_graph(&pg.second); |
| 4227 | } |
| 4228 | } |
| 4229 | |
| 4230 | return Status::OK(); |
| 4231 | } |
| 4232 | |
| 4233 | } // namespace tensorflow |
| 4234 | |