| 248 | } |
| 249 | |
| 250 | Status XRTTupleAllocation::WriteLiteral(xla::Backend* backend, |
| 251 | const xla::Literal& literal) { |
| 252 | if (!xla::ShapeUtil::Equal(literal.shape(), on_host_shape())) { |
| 253 | return errors::InvalidArgument( |
| 254 | "New literal shape not matching the existing one: literal=", |
| 255 | xla::ShapeUtil::HumanStringWithLayout(literal.shape()), |
| 256 | " device=", xla::ShapeUtil::HumanStringWithLayout(on_host_shape())); |
| 257 | } |
| 258 | mutex_lock lock(lock_); |
| 259 | if (literal_ != nullptr) { |
| 260 | // The allocation is currently swapped out, and we have a host literal for |
| 261 | // its content. Just update the host literal with the new value. |
| 262 | return literal_->CopyFrom(literal); |
| 263 | } |
| 264 | TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, ToShapedBuffer()); |
| 265 | auto transfer_manager = backend->transfer_manager(); |
| 266 | TF_ASSIGN_OR_RETURN(auto stream, backend->BorrowStream(device_ordinal())); |
| 267 | return transfer_manager->TransferLiteralToDevice(stream.get(), literal, |
| 268 | shaped_buffer); |
| 269 | } |
| 270 | |
| 271 | xla::StatusOr<bool> XRTTupleAllocation::SwapOut(xla::Backend* backend, |
| 272 | bool swap_pinned) { |
no test coverage detected