| 29 | outfeed_slices_(std::move(outfeed_slices)) {} |
| 30 | |
| 31 | Status OutfeedThunk::ExecuteOnStream(const ExecuteParams& params) { |
| 32 | auto& stream = *params.stream; |
| 33 | auto& buffer_allocations = *params.buffer_allocations; |
| 34 | |
| 35 | VLOG(2) << "Outfeeding from GPU: " << hlo_instruction()->ToString(); |
| 36 | |
| 37 | auto op_profiler = |
| 38 | params.profiler->MakeScopedInstructionProfiler(hlo_instruction()); |
| 39 | OutfeedManager* outfeed_manager = GetOrCreateOutfeedManager(); |
| 40 | ShapeTree<std::unique_ptr<OutfeedBuffer>>* outfeed_buffers = |
| 41 | outfeed_manager->BlockingGetNextDestination(); |
| 42 | |
| 43 | // Nothing to be done for empty tuples. |
| 44 | if (ShapeUtil::IsEmptyTuple(hlo_instruction()->operand(0)->shape())) { |
| 45 | return Status::OK(); |
| 46 | } |
| 47 | CHECK(ShapeUtil::Compatible(hlo_instruction()->operand(0)->shape(), |
| 48 | outfeed_buffers->shape())); |
| 49 | |
| 50 | TF_RETURN_IF_ERROR(outfeed_buffers->ForEachMutableElementWithStatus( |
| 51 | [&](const ShapeIndex& index, std::unique_ptr<OutfeedBuffer>* buffer) { |
| 52 | if (!*buffer) { // Tuple pointers. |
| 53 | return Status::OK(); |
| 54 | } |
| 55 | |
| 56 | BufferAllocation::Slice slice = outfeed_slices_.element(index); |
| 57 | se::DeviceMemoryBase data_address; |
| 58 | if (slice.allocation()) { |
| 59 | // If we have a static allocation, read it from there. This avoids |
| 60 | // synchronizing the host and device just to read a pointer. |
| 61 | data_address = buffer_allocations.GetDeviceAddress(slice); |
| 62 | } else { |
| 63 | // Otherwise we have to read the tuple pointer first. |
| 64 | CHECK(!index.empty()); |
| 65 | // Copy the parent buffer to the host. |
| 66 | BufferAllocation::Slice tuple_slice = |
| 67 | outfeed_slices_.element(ShapeIndexView(index).ConsumeFront()); |
| 68 | if (!tuple_slice.allocation()) { |
| 69 | return Unimplemented( |
| 70 | "Nested dynamic tuples are not supported on GPU"); |
| 71 | } |
| 72 | se::DeviceMemoryBase tuple_address = |
| 73 | buffer_allocations.GetDeviceAddress(tuple_slice); |
| 74 | CHECK(tuple_slice.size() % sizeof(void*) == 0) |
| 75 | << "Tuple size must be a multiple of pointer size"; |
| 76 | std::vector<void*> tuple_element_buffer_addresses(tuple_slice.size() / |
| 77 | sizeof(void*)); |
| 78 | stream.ThenMemcpy(tuple_element_buffer_addresses.data(), |
| 79 | tuple_address, tuple_slice.size()); |
| 80 | TF_RETURN_IF_ERROR(stream.BlockHostUntilDone()); |
| 81 | // The data address is specified by the element of the tuple pointer |
| 82 | // buffer. |
| 83 | data_address = |
| 84 | se::DeviceMemoryBase(tuple_element_buffer_addresses[index.back()], |
| 85 | (*buffer)->length()); |
| 86 | } |
| 87 | |
| 88 | // TODO(b/111309141): Run this on a separate stream so it doesn't block |
nothing calls this directly
no test coverage detected