| 317 | } |
| 318 | |
| 319 | xla::StatusOr<jit::DeviceId> InferDeviceForCluster( |
| 320 | jit::DeviceInfoCache* device_info_cache, Node* n, |
| 321 | const string& function_name, const FunctionLibraryDefinition& flib_def) { |
| 322 | const FunctionDef* func_def = flib_def.Find(function_name); |
| 323 | TF_RET_CHECK(func_def) << "Could not find " << function_name; |
| 324 | |
| 325 | jit::DeviceSet device_set; |
| 326 | |
| 327 | for (const NodeDef& ndef : func_def->node_def()) { |
| 328 | VLOG(3) << ndef.DebugString(); |
| 329 | if (!ndef.device().empty()) { |
| 330 | TF_ASSIGN_OR_RETURN(jit::DeviceId device_id, |
| 331 | device_info_cache->GetIdFor(ndef.device())); |
| 332 | device_set.Insert(device_id); |
| 333 | } |
| 334 | } |
| 335 | |
| 336 | if (!n->assigned_device_name().empty()) { |
| 337 | // TODO(sanjoy): We need this because EncapsulateSubgraphsPass drops device |
| 338 | // assignment when constant folding. We should fix EncapsulateSubgraphsPass |
| 339 | // instead. |
| 340 | TF_ASSIGN_OR_RETURN(jit::DeviceId device_id, |
| 341 | device_info_cache->GetIdFor(n->assigned_device_name())); |
| 342 | device_set.Insert(device_id); |
| 343 | } |
| 344 | |
| 345 | TF_ASSIGN_OR_RETURN(jit::DeviceId result, |
| 346 | PickDeviceForXla(*device_info_cache, device_set, |
| 347 | /*allow_mixing_unknown_and_cpu=*/true)); |
| 348 | VLOG(2) << "For " << function_name << " PickDeviceForXla(" |
| 349 | << device_info_cache->DebugString(device_set) << ") -> " |
| 350 | << device_info_cache->GetNameFor(result); |
| 351 | return result; |
| 352 | } |
| 353 | |
| 354 | std::vector<Output> GetXlaRunArgs(const Scope& s, |
| 355 | const XlaClusterInfo& cluster_info, |
nothing calls this directly
no test coverage detected