| 323 | } |
| 324 | |
| 325 | xla::StatusOr<std::pair<XlaDeviceContext*, XlaDeviceContext*>> |
| 326 | XlaDevice::GetDeviceContextLocked() { |
| 327 | xla::Backend* backend = client()->mutable_backend(); |
| 328 | |
| 329 | // Ensure all our streams are valid, borrowing new streams if necessary. |
| 330 | bool need_new_device_context = !device_context_; |
| 331 | TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_, |
| 332 | &need_new_device_context)); |
| 333 | |
| 334 | std::shared_ptr<se::Stream> host_to_device_stream; |
| 335 | std::shared_ptr<se::Stream> device_to_host_stream; |
| 336 | std::vector<std::shared_ptr<se::Stream>> device_to_device_streams; |
| 337 | if (use_multiple_streams_) { |
| 338 | TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "host_to_device_stream", |
| 339 | &host_to_device_stream_, |
| 340 | &need_new_device_context)); |
| 341 | for (std::shared_ptr<se::Stream>& stream : device_to_device_streams_) { |
| 342 | TF_RETURN_IF_ERROR( |
| 343 | EnsureStreamOkLocked(backend, "device_to_device_stream", &stream, |
| 344 | &need_new_device_context)); |
| 345 | } |
| 346 | host_to_device_stream = host_to_device_stream_; |
| 347 | device_to_device_streams = device_to_device_streams_; |
| 348 | // The data transfer requests from device to host could arrive out of order, |
| 349 | // so a single stream would cause deadlock. For this case, |
| 350 | // xla_device_context would borrow a stream for each transfer request. |
| 351 | device_to_host_stream = nullptr; |
| 352 | } else { |
| 353 | host_to_device_stream = stream_; |
| 354 | device_to_host_stream = stream_; |
| 355 | device_to_device_streams = {stream_}; |
| 356 | } |
| 357 | |
| 358 | if (!need_new_device_context) { |
| 359 | return std::make_pair(device_context_, fast_mem_device_context_); |
| 360 | } |
| 361 | |
| 362 | // At this point we know we need a new device context. |
| 363 | // Call GetAllocator for the side-effect of ensuring the allocator is created. |
| 364 | GetAllocatorLocked({}); |
| 365 | if (device_context_) { |
| 366 | device_context_->Unref(); |
| 367 | } |
| 368 | if (fast_mem_device_context_) { |
| 369 | fast_mem_device_context_->Unref(); |
| 370 | } |
| 371 | // The XlaDeviceContext keeps a reference count to the streams, and the |
| 372 | // XlaDeviceContext remains live for the duration of a Executor run. This |
| 373 | // ensures that the streams remain live for the duration of a run, even if |
| 374 | // an error is encountered and the streams are replaced with new ones. |
| 375 | device_context_ = new XlaDeviceContext( |
| 376 | stream_, host_to_device_stream, device_to_host_stream, |
| 377 | device_to_device_streams, client(), shape_representation_fn_, |
| 378 | thread_pool_.get(), false); |
| 379 | VLOG(1) << "XlaDevice " << this << " new XlaDeviceContext(fast_mem=false) " |
| 380 | << device_context_; |
| 381 | |
| 382 | fast_mem_device_context_ = new XlaDeviceContext( |
nothing calls this directly
no test coverage detected