| 260 | } // namespace |
| 261 | |
| 262 | bool IsSingleGpuGraph(const Graph& g) { |
| 263 | int gpus_seen = 0; |
| 264 | absl::flat_hash_set<string> devices_seen; |
| 265 | |
| 266 | for (Node* n : g.op_nodes()) { |
| 267 | if (devices_seen.contains(n->assigned_device_name())) { |
| 268 | continue; |
| 269 | } |
| 270 | |
| 271 | int gpu_number = GetGpuNumber(n->assigned_device_name()); |
| 272 | if (gpu_number != -1) { |
| 273 | if (++gpus_seen > 1) { |
| 274 | return false; |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | devices_seen.insert(n->assigned_device_name()); |
| 279 | } |
| 280 | |
| 281 | return gpus_seen == 1; |
| 282 | } |
| 283 | |
| 284 | OptimizerOptions::GlobalJitLevel GetGlobalJitLevelForGraph( |
| 285 | const GraphOptimizationPassOptions& options) { |