| 280 | } |
| 281 | |
| 282 | /* static */ StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::MakeTuple( |
| 283 | const std::vector<PyLocalBuffer*> buffers, |
| 284 | std::shared_ptr<PyLocalClient> client, std::shared_ptr<Device> device) { |
| 285 | TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device, |
| 286 | device->GetLocalDeviceState()); |
| 287 | std::vector<Shape> host_shapes; |
| 288 | std::vector<Shape> device_shapes; |
| 289 | std::vector<std::shared_ptr<SharedDeviceBuffer>> device_buffers; |
| 290 | host_shapes.reserve(buffers.size()); |
| 291 | device_shapes.reserve(buffers.size()); |
| 292 | device_buffers.reserve(buffers.size()); |
| 293 | for (const PyLocalBuffer* buffer : buffers) { |
| 294 | if (buffer->device().get() != device.get()) { |
| 295 | return InvalidArgument( |
| 296 | "Tuple elements must be on the same device; %s vs %s", |
| 297 | buffer->device()->DebugString(), device->DebugString()); |
| 298 | } |
| 299 | std::shared_ptr<SharedDeviceBuffer> device_buffer = buffer->DeviceBuffer(); |
| 300 | if (!device_buffer) { |
| 301 | return InvalidArgument( |
| 302 | "Invalid buffer passed to MakeTuple() as argument %d.", |
| 303 | device_buffers.size()); |
| 304 | } |
| 305 | host_shapes.push_back(buffer->on_host_shape()); |
| 306 | device_shapes.push_back(buffer->on_device_shape()); |
| 307 | device_buffers.push_back(std::move(device_buffer)); |
| 308 | } |
| 309 | se::DeviceMemoryAllocator* allocator = client->allocator(); |
| 310 | TransferManager* transfer_manager = |
| 311 | client->client()->backend().transfer_manager(); |
| 312 | |
| 313 | Shape on_host_shape = ShapeUtil::MakeTupleShape(host_shapes); |
| 314 | auto definition_event = std::make_shared<BufferDefinitionEvent>(); |
| 315 | TF_ASSIGN_OR_RETURN( |
| 316 | std::shared_ptr<SharedDeviceBuffer> tuple_buffer, |
| 317 | SharedDeviceBuffer::MakeTuple( |
| 318 | device_buffers, on_host_shape, transfer_manager, allocator, |
| 319 | local_device->device_ordinal(), definition_event)); |
| 320 | auto buffer = absl::make_unique<PyLocalBuffer>( |
| 321 | std::move(on_host_shape), ShapeUtil::MakeTupleShape(device_shapes), |
| 322 | tuple_buffer, std::move(client), std::move(device)); |
| 323 | |
| 324 | // TODO(phawkins): extend TransferManager so we do not need to form a full |
| 325 | // ShapedBuffer just to write the root tuple index table. |
| 326 | TF_ASSIGN_OR_RETURN(ShapedBuffer shaped_buffer, buffer->AsShapedBuffer()); |
| 327 | if (!transfer_manager->CanShapedBufferBeAccessedNow( |
| 328 | local_device->host_to_device_stream()->parent(), shaped_buffer)) { |
| 329 | // Wait for the compute stream so that memory allocations are synchronized. |
| 330 | local_device->host_to_device_stream()->ThenWaitFor( |
| 331 | local_device->compute_stream()); |
| 332 | } |
| 333 | TF_RETURN_IF_ERROR(transfer_manager->WriteRootTupleIndexTable( |
| 334 | local_device->host_to_device_stream(), shaped_buffer)); |
| 335 | |
| 336 | TF_ASSIGN_OR_RETURN(EventPool::Handle event, |
| 337 | local_device->event_pool().ThenAllocateAndRecordEvent( |
| 338 | local_device->host_to_device_stream())); |
| 339 | definition_event->SetDefinitionEvent(std::move(event), |
nothing calls this directly
no test coverage detected