| 282 | } |
| 283 | |
| 284 | xla::StatusOr<bool> XRTTupleAllocation::SwapIn(XRTMemoryManager* memory_manager, |
| 285 | xla::Backend* backend) { |
| 286 | // We need to call AllocateScopedShapedBuffer() outside the locks, since the |
| 287 | // XRTMemoryManager might end up calling back into the SwapOut() API. |
| 288 | // So we do a quick check before using the IsSwapped() API, and it can happen |
| 289 | // that the allocation becomes swapped in after the check. This means which we |
| 290 | // will end up doing an allocation, and then releasing it soon after (via its |
| 291 | // scoped variables). This is an unlikely scenario (two threads calling |
| 292 | // SwapIn() on the same allocation) though. |
| 293 | if (!IsSwapped()) { |
| 294 | return false; |
| 295 | } |
| 296 | |
| 297 | auto transfer_manager = backend->transfer_manager(); |
| 298 | std::unique_ptr<xla::ScopedShapedBuffer> scoped_buffer; |
| 299 | TF_RETURN_IF_ERROR( |
| 300 | AllocateScopedShapedBuffer(memory_manager, backend, device_ordinal(), |
| 301 | on_host_shape(), &scoped_buffer)); |
| 302 | TF_ASSIGN_OR_RETURN(auto stream, backend->BorrowStream(device_ordinal())); |
| 303 | |
| 304 | mutex_lock lock(lock_); |
| 305 | if (literal_ != nullptr) { |
| 306 | TF_RETURN_IF_ERROR(transfer_manager->TransferLiteralToDevice( |
| 307 | stream.get(), *literal_, *scoped_buffer)); |
| 308 | |
| 309 | auto shaped_buffer = scoped_buffer->release(); |
| 310 | InitializeFromShapedBuffer(shaped_buffer, backend->memory_allocator(), |
| 311 | device_ordinal()); |
| 312 | literal_ = nullptr; |
| 313 | return true; |
| 314 | } |
| 315 | return false; |
| 316 | } |
| 317 | |
| 318 | xla::StatusOr<bool> XRTTupleAllocation::PinAndSwapIn( |
| 319 | XRTMemoryManager* memory_manager, xla::Backend* backend) { |
no test coverage detected