static */
| 162 | } |
| 163 | |
| 164 | /* static */ StatusOr<std::vector<se::StreamExecutor*>> |
| 165 | PlatformUtil::GetStreamExecutors( |
| 166 | se::Platform* platform, |
| 167 | const absl::optional<std::set<int>>& allowed_devices) { |
| 168 | int device_count = platform->VisibleDeviceCount(); |
| 169 | if (device_count <= 0) { |
| 170 | return NotFound("no %s devices found", platform->Name()); |
| 171 | } |
| 172 | if (platform->id() == se::host::kHostPlatformId) { |
| 173 | // On host "devices", StreamExecutor exports a device for each hardware |
| 174 | // thread. Because we parallelize a single computation across threads, it |
| 175 | // doesn't make sense to expose these as separate devices, so by default we |
| 176 | // fix the number of devices to one. However we do let the user override |
| 177 | // this behavior to help run tests on the host that run models in parallel |
| 178 | // across multiple devices. |
| 179 | device_count = |
| 180 | GetDebugOptionsFromFlags().xla_force_host_platform_device_count(); |
| 181 | } |
| 182 | std::vector<se::StreamExecutor*> stream_executors(device_count, nullptr); |
| 183 | VLOG(1) << "Initializing devices"; |
| 184 | { |
| 185 | tensorflow::thread::ThreadPool thread_pool( |
| 186 | tensorflow::Env::Default(), "device_initialization", device_count); |
| 187 | for (int i = 0; i < device_count; ++i) { |
| 188 | // Once a stream executor is instantiated it will cause allocations on |
| 189 | // the device, for example for GPUs cuda context, cudnn handles etc. will |
| 190 | // be constructed. By constructing stream executors only on the |
| 191 | // allowed_devices, we don't make any allocations on other devices. |
| 192 | // This helps in multi-process executions on the same host like horovod or |
| 193 | // shared hosts. |
| 194 | if (allowed_devices && allowed_devices->count(i) == 0) { |
| 195 | VLOG(1) << "Not initializing StreamExecutor for device " << i |
| 196 | << " since it is not in the visible device list"; |
| 197 | continue; |
| 198 | } |
| 199 | thread_pool.Schedule([platform, i, &stream_executors]() { |
| 200 | VLOG(1) << "Started device init " << i; |
| 201 | se::StreamExecutorConfig config; |
| 202 | config.ordinal = i; |
| 203 | auto executor_status = platform->GetExecutor(config); |
| 204 | if (executor_status.ok()) { |
| 205 | se::StreamExecutor* executor = executor_status.ValueOrDie(); |
| 206 | if (IsDeviceSupported(executor)) { |
| 207 | stream_executors[i] = executor; |
| 208 | } |
| 209 | } else { |
| 210 | LOG(WARNING) << "unable to create StreamExecutor for " |
| 211 | << platform->Name() << ":" << i << ": " |
| 212 | << executor_status.status().error_message(); |
| 213 | } |
| 214 | VLOG(1) << "Finished device init " << i; |
| 215 | }); |
| 216 | } |
| 217 | // Block here in thread_pool destructor until all devices are initialized. |
| 218 | } |
| 219 | VLOG(1) << "Device initialization complete"; |
| 220 | |
| 221 | std::vector<se::StreamExecutor*> out; |
nothing calls this directly
no test coverage detected