| 592 | } |
| 593 | |
| 594 | Status XRTTupleAllocation::AliasBufferFrom(const XRTTupleAllocation& source, |
| 595 | const xla::ShapeIndex& source_index, |
| 596 | const xla::ShapeIndex& dest_index) { |
| 597 | XRTBufferAllocation* source_buffer = source.buffers_.element(source_index); |
| 598 | XRTBufferAllocation* dest_buffer = buffers_.element(dest_index); |
| 599 | if (dest_buffer != nullptr) { |
| 600 | // We allow the destination size being zero, because there are cases where |
| 601 | // we are coming in later filling in null/uninitialized device buffers. In |
| 602 | // all other cases, the size of the new buffer must match. |
| 603 | if (source_buffer->allocation().size() != |
| 604 | dest_buffer->allocation().size() && |
| 605 | dest_buffer->allocation().size() != 0) { |
| 606 | return errors::InvalidArgument( |
| 607 | "Source buffer at index ", source_index.ToString(), |
| 608 | " does not match the size of destination buffer at index ", |
| 609 | dest_index.ToString(), ": ", source_buffer->allocation().size(), |
| 610 | " vs ", dest_buffer->allocation().size()); |
| 611 | } |
| 612 | } else { |
| 613 | const xla::Shape& source_subshape = |
| 614 | xla::ShapeUtil::GetSubshape(source.on_device_shape(), source_index); |
| 615 | const xla::Shape& dest_subshape = |
| 616 | xla::ShapeUtil::GetSubshape(on_device_shape(), dest_index); |
| 617 | if (!xla::ShapeUtil::Equal(source_subshape, dest_subshape)) { |
| 618 | return errors::InvalidArgument( |
| 619 | "Source and destination subshapes do not match: source=", |
| 620 | xla::ShapeUtil::HumanStringWithLayout(source_subshape), |
| 621 | " dest=", xla::ShapeUtil::HumanStringWithLayout(dest_subshape)); |
| 622 | } |
| 623 | } |
| 624 | *buffers_.mutable_element(dest_index) = source_buffer; |
| 625 | source_buffer->Ref(); |
| 626 | if (dest_buffer != nullptr) { |
| 627 | // If we handed over the ownership of a buffer in ToDeviceMemoryTree(), we |
| 628 | // will be called here on the way back from execution, to alias back the |
| 629 | // buffer at that index. In that case the buffers will be the same. So we |
| 630 | // need to discard the memory at the destination buffer, before releasing |
| 631 | // the reference. |
| 632 | if (dest_buffer->allocation().IsSameAs(source_buffer->allocation()) && |
| 633 | dest_buffer != source_buffer) { |
| 634 | dest_buffer->DiscardAllocation(); |
| 635 | } |
| 636 | dest_buffer->Unref(); |
| 637 | } |
| 638 | return Status::OK(); |
| 639 | } |
| 640 | |
| 641 | xla::StatusOr<xla::ShapeTree<xla::MaybeOwningDeviceMemory>> |
| 642 | XRTTupleAllocation::ToDeviceMemoryTree( |
no test coverage detected