| 98 | } |
| 99 | |
| 100 | xla::StatusOr<absl::optional<jit::DeviceId>> PickDeviceForXlaImpl( |
| 101 | const jit::DeviceInfoCache& device_info_cache, |
| 102 | const jit::DeviceSet& devices, bool allow_mixing_unknown_and_cpu, |
| 103 | bool failure_to_pick_is_error) { |
| 104 | #define FAILED_TO_PICK_DEVICE(failing_status) \ |
| 105 | do { \ |
| 106 | if (failure_to_pick_is_error) { \ |
| 107 | return failing_status; \ |
| 108 | } else { \ |
| 109 | return {absl::nullopt}; \ |
| 110 | } \ |
| 111 | } while (false) |
| 112 | |
| 113 | absl::optional<jit::DeviceId> maybe_gpu_device; |
| 114 | absl::optional<jit::DeviceId> maybe_cpu_device; |
| 115 | absl::optional<jit::DeviceId> maybe_unknown_device; |
| 116 | |
| 117 | bool multiple_cpu_devices = false; |
| 118 | bool multiple_gpu_devices = false; |
| 119 | bool multiple_unknown_devices = false; |
| 120 | |
| 121 | // Returns 'true' if d0 and d1 are conflicting devices. If they are |
| 122 | // compatible, update d1 with a more specific one. |
| 123 | // TODO(sanjoy): Cache DeviceNameUtils::ParsedName inside device_info_cache. |
| 124 | const auto is_multiple_devices = |
| 125 | [&](const jit::DeviceId& d0, absl::optional<jit::DeviceId>* d1) -> bool { |
| 126 | const absl::string_view name0 = device_info_cache.GetNameFor(d0); |
| 127 | const absl::string_view name1 = device_info_cache.GetNameFor(d1->value()); |
| 128 | |
| 129 | DeviceNameUtils::ParsedName parsed0, parsed1; |
| 130 | if (!DeviceNameUtils::ParseFullName(name0, &parsed0) || |
| 131 | !DeviceNameUtils::ParseFullName(name1, &parsed1) || |
| 132 | !DeviceNameUtils::AreCompatibleDevNames(parsed0, parsed1)) { |
| 133 | return true; |
| 134 | } |
| 135 | |
| 136 | if (DeviceNameUtils::IsSpecification(parsed0, parsed1)) { |
| 137 | return false; |
| 138 | } |
| 139 | |
| 140 | if (DeviceNameUtils::IsSpecification(parsed1, parsed0)) { |
| 141 | *d1 = d0; |
| 142 | return false; |
| 143 | } |
| 144 | |
| 145 | return true; |
| 146 | }; |
| 147 | |
| 148 | devices.ForEach([&](jit::DeviceId device) { |
| 149 | if (device_info_cache.IsGpu(device)) { |
| 150 | if (maybe_gpu_device) { |
| 151 | multiple_gpu_devices = is_multiple_devices(device, &maybe_gpu_device); |
| 152 | if (multiple_gpu_devices) return false; |
| 153 | } else { |
| 154 | maybe_gpu_device = device; |
| 155 | } |
| 156 | } else if (device_info_cache.IsCpu(device)) { |
| 157 | if (maybe_cpu_device) { |
no test coverage detected