| 44 | } |
| 45 | |
| 46 | Status XlaTensor::AllocateShapedBuffer(DataType dtype, |
| 47 | const xla::Shape& on_host_shape, |
| 48 | xla::LocalClient* client, |
| 49 | int device_ordinal) { |
| 50 | xla::Shape on_device_shape = |
| 51 | client->backend().transfer_manager()->HostShapeToDeviceShape( |
| 52 | on_host_shape); |
| 53 | |
| 54 | xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape, |
| 55 | client->backend().memory_allocator(), |
| 56 | device_ordinal); |
| 57 | for (auto& index_to_buffer : shaped_buffer.buffers()) { |
| 58 | xla::Shape subshape = |
| 59 | xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); |
| 60 | uint64 size = |
| 61 | client->backend().transfer_manager()->GetByteSizeRequirement(subshape); |
| 62 | TF_ASSIGN_OR_RETURN(se::OwningDeviceMemory buffer, |
| 63 | client->backend().memory_allocator()->Allocate( |
| 64 | device_ordinal, size, /*retry_on_failure=*/false, |
| 65 | subshape.layout().memory_space())); |
| 66 | // Move our buffer into shaped_buffer, which takes ownership of it. |
| 67 | index_to_buffer.second = buffer.Release(); |
| 68 | } |
| 69 | |
| 70 | VLOG(4) << shaped_buffer.ToString(); |
| 71 | |
| 72 | set_shaped_buffer(std::move(shaped_buffer)); |
| 73 | return Status::OK(); |
| 74 | } |
| 75 | |
| 76 | void XlaTensor::WaitForDefinitionEventOnStream(se::Stream* stream) { |
| 77 | mutex_lock lock(mu_); |
no test coverage detected