| 38 | } |
| 39 | |
| 40 | StatusOr<std::unique_ptr<BufferAllocations>> BufferAllocations::Builder::Build( |
| 41 | const BufferAssignment* buffer_assignment, int device_ordinal, |
| 42 | se::DeviceMemoryAllocator* memory_allocator) { |
| 43 | const int64 num_buffers = buffer_assignment->Allocations().size(); |
| 44 | auto buffer_allocations = absl::WrapUnique(new BufferAllocations( |
| 45 | num_buffers, device_ordinal, memory_allocator, buffer_assignment)); |
| 46 | |
| 47 | for (BufferAllocation::Index i = 0; i < num_buffers; ++i) { |
| 48 | const BufferAllocation& allocation = buffer_assignment->GetAllocation(i); |
| 49 | const int64 expected_alignment = [&] { |
| 50 | if (allocation.is_entry_computation_parameter()) { |
| 51 | return kEntryParameterAlignBytes; |
| 52 | } else if (allocation.is_constant()) { |
| 53 | return kConstantBufferAlignBytes; |
| 54 | } else { |
| 55 | return kXlaAllocatedBufferAlignBytes; |
| 56 | } |
| 57 | }(); |
| 58 | |
| 59 | // If buffer #i's address is already registered (e.g. external arguments or |
| 60 | // result buffers), use that registered buffer. |
| 61 | if (se::DeviceMemoryBase* address = |
| 62 | tensorflow::gtl::FindOrNull(registered_buffers_, i)) { |
| 63 | if (reinterpret_cast<uintptr_t>(address->opaque()) % expected_alignment != |
| 64 | 0) { |
| 65 | return InternalError( |
| 66 | "Address of registered buffer %d must be a multiple of %x, but " |
| 67 | "was %p", |
| 68 | i, kEntryParameterAlignBytes, address->opaque()); |
| 69 | } |
| 70 | buffer_allocations->SetBuffer(i, *address); |
| 71 | continue; |
| 72 | } |
| 73 | |
| 74 | // Allocate each allocation that might escape, or is the temp buffer. |
| 75 | bool seen_temp_buffer = false; |
| 76 | if (allocation.maybe_live_out() || allocation.IsPreallocatedTempBuffer()) { |
| 77 | const int64 buffer_size = allocation.size(); |
| 78 | se::DeviceMemoryBase buffer_address; |
| 79 | if (buffer_size > 0) { |
| 80 | se::OwningDeviceMemory buffer; |
| 81 | TF_ASSIGN_OR_RETURN( |
| 82 | buffer, memory_allocator->Allocate(device_ordinal, buffer_size)); |
| 83 | if (reinterpret_cast<uintptr_t>(buffer->opaque()) % |
| 84 | expected_alignment != |
| 85 | 0) { |
| 86 | return InternalError( |
| 87 | "Address returned by memory_allocator->Allocate must be a " |
| 88 | "multiple of 0x%x, but was %p", |
| 89 | kXlaAllocatedBufferAlignBytes, buffer->opaque()); |
| 90 | } |
| 91 | // We do manual memory management within BufferAllocations. Be sure not |
| 92 | // to do a TF_RETURN_IF_ERROR between this line and the |
| 93 | // buffer_allocations->SetBuffer(buffer_address) call below! |
| 94 | buffer_address = buffer.Release(); |
| 95 | } |
| 96 | |
| 97 | buffer_allocations->SetBuffer(i, buffer_address); |
nothing calls this directly
no test coverage detected