| 49 | } |
| 50 | |
| 51 | xla::StatusOr<DeviceId> DeviceInfoCache::GetIdFor(absl::string_view name) { |
| 52 | TF_RET_CHECK(!name.empty()); |
| 53 | |
| 54 | auto it = name_to_id_.find(name); |
| 55 | if (it != name_to_id_.end()) { |
| 56 | return it->second; |
| 57 | } |
| 58 | |
| 59 | int new_id = names_.size(); |
| 60 | names_.push_back(string(name)); |
| 61 | id_to_device_type_.push_back(absl::make_unique<DeviceType>("")); |
| 62 | DeviceType* device_type = id_to_device_type_.back().get(); |
| 63 | TF_RETURN_IF_ERROR(DeviceNameToDeviceType(names_.back(), device_type)); |
| 64 | |
| 65 | is_cpu_.push_back(device_type->type_string() == DEVICE_CPU); |
| 66 | is_gpu_.push_back(device_type->type_string() == DEVICE_GPU); |
| 67 | |
| 68 | name_to_id_.emplace(string(name), DeviceId(new_id)); |
| 69 | |
| 70 | const XlaOpRegistry::DeviceRegistration* compilation_device; |
| 71 | if (!XlaOpRegistry::GetCompilationDevice(device_type->type(), |
| 72 | &compilation_device)) { |
| 73 | compilation_device = nullptr; |
| 74 | } |
| 75 | id_to_compilation_device_.push_back(compilation_device); |
| 76 | |
| 77 | return DeviceId(new_id); |
| 78 | } |
| 79 | |
| 80 | string DeviceInfoCache::DebugString(const DeviceSet& device_set) const { |
| 81 | std::vector<string> names; |