Method MakeSubBuffer

tensorflow/compiler/xrt/xrt_state.cc:354–393

static*/

}

/*static*/ Status XRTTupleAllocation::MakeSubBuffer(
    XRTTupleAllocation* parent, const xla::ShapeIndex& subshape,
    XRTTupleAllocation** allocation, bool alias_parent_allocation) {
  TF_ASSIGN_OR_RETURN(
      const xla::Shape* host_sub_shape,
      xla::ShapeUtil::TryGetSubshape(parent->on_host_shape(), subshape));
  TF_ASSIGN_OR_RETURN(
      const xla::Shape* device_sub_shape,
      xla::ShapeUtil::TryGetSubshape(parent->on_device_shape(), subshape));

  *allocation =
      new XRTTupleAllocation(parent->device_ordinal(), parent->allocator_,
                             *host_sub_shape, *device_sub_shape);
  if (alias_parent_allocation) {
    // Copy the subtree of allocations from the parent allocation.
    (*allocation)->buffers_.CopySubtreeFrom(parent->buffers_, subshape, {});
    // Increment the refcount on each aliased buffer.
    (*allocation)
        ->buffers_.ForEachElement(
            [](const xla::ShapeIndex& index,
               const XRTBufferAllocationPtr& buffer) { buffer->Ref(); });
  } else {
    // Find the buffers in the parent allocation that match the subtree, and
    // move the parent allocation's buffer over to the new allocation.
    (*allocation)
        ->buffers_.ForEachMutableElement(
            [&](const xla::ShapeIndex& index, XRTBufferAllocationPtr* buffer) {
              // Extend the allocation's index to the parent's frame by adding
              // subshape as a prefix.
              xla::ShapeIndex parent_index = subshape;
              for (int i = 0; i < index.size(); ++i) {
                parent_index.push_back(index[i]);
              }
              *buffer = parent->buffers_.element(parent_index);
              *parent->buffers_.mutable_element(parent_index) = nullptr;
            });
  }
  (*allocation)->SetDeviceMemorySize();
  return Status::OK();
}

void XRTTupleAllocation::SetDeviceMemorySize() {
  size_t size = 0;