static */
| 403 | } |
| 404 | |
| 405 | /* static */ Status XRTTupleAllocation::ExpandTreeOfTuples( |
| 406 | const xla::ShapeTree<ExpandedTupleInput>& elements, int device_ordinal, |
| 407 | se::DeviceMemoryAllocator* allocator, xla::Shape* host_shape, |
| 408 | xla::Shape* device_shape) { |
| 409 | // Initialize both host and device shape to be the 'spine' of the new tuple |
| 410 | // shape, given by the shape of the tree of tuples. |
| 411 | *host_shape = elements.shape(); |
| 412 | *device_shape = elements.shape(); |
| 413 | // Now go over the leaves of the tree of tuples, and 'graft' the host/device |
| 414 | // shapes of the allocation at that leaf onto the expanded host/device shapes |
| 415 | // at the leaf position. |
| 416 | TF_RETURN_IF_ERROR(elements.ForEachElementWithStatus( |
| 417 | [&](const xla::ShapeIndex& index, const ExpandedTupleInput& element) { |
| 418 | if (elements.IsLeaf(index)) { |
| 419 | if (element.allocation == nullptr) { |
| 420 | return errors::InvalidArgument( |
| 421 | "MakeTuple elements has a null internal node at index ", |
| 422 | index.ToString()); |
| 423 | } |
| 424 | if (device_ordinal != element.allocation->device_ordinal() || |
| 425 | allocator != element.allocation->allocator_) { |
| 426 | return errors::InvalidArgument( |
| 427 | "MakeTuple elements must all be allocated on the same device " |
| 428 | "as the destination."); |
| 429 | } |
| 430 | *xla::ShapeUtil::GetMutableSubshape(host_shape, index) = |
| 431 | element.allocation->on_host_shape(); |
| 432 | *xla::ShapeUtil::GetMutableSubshape(device_shape, index) = |
| 433 | element.allocation->on_device_shape(); |
| 434 | } else { |
| 435 | if (element.allocation != nullptr) { |
| 436 | return errors::InvalidArgument( |
| 437 | "MakeTuple elements has a non-null internal node at index ", |
| 438 | index.ToString()); |
| 439 | } |
| 440 | } |
| 441 | return Status::OK(); |
| 442 | })); |
| 443 | return Status::OK(); |
| 444 | } |
| 445 | |
| 446 | /*static*/ Status XRTTupleAllocation::MakeTuple( |
| 447 | XRTMemoryManager* memory_manager, xla::Backend* backend, int device_ordinal, |
nothing calls this directly
no test coverage detected