static */
| 156 | |
| 157 | /* static */ |
| 158 | StatusOr<std::unique_ptr<PyLocalBuffer>> PyLocalBuffer::FromHostBuffer( |
| 159 | const void* data, const Shape& shape, bool force_copy, |
| 160 | std::shared_ptr<void> buffer_reference, |
| 161 | std::shared_ptr<PyLocalClient> client, std::shared_ptr<Device> device) { |
| 162 | tensorflow::profiler::TraceMe traceme("PyLocalBuffer::FromLiterals"); |
| 163 | VLOG(2) << "PyLocalBuffer::FromLiterals: shape: " << shape.ToString() |
| 164 | << " device: " << device->DebugString(); |
| 165 | TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device, |
| 166 | device->GetLocalDeviceState()); |
| 167 | |
| 168 | // If we are on the host platform and the input buffer is sufficiently |
| 169 | // aligned, we can simply point to the NumPy array's data without any further |
| 170 | // copies. We require a 64-byte alignment because XLA may generate AVX512 |
| 171 | // code which requires it. Unfortunately NumPy's allocator doesn't align |
| 172 | // quite as aggressively, so there's a high chance this test will fail. |
| 173 | static constexpr int kMinimumAlignment = 64; |
| 174 | if (!force_copy && |
| 175 | ((absl::bit_cast<std::uintptr_t>(data) & (kMinimumAlignment - 1)) == 0) && |
| 176 | local_device->executor()->platform_kind() == se::PlatformKind::kHost) { |
| 177 | std::function<void()> on_delete_callback = |
| 178 | [buffer_reference{std::move(buffer_reference)}]() { |
| 179 | // Frees buffer_reference. |
| 180 | }; |
| 181 | se::DeviceMemoryBase buffer(const_cast<void*>(data), |
| 182 | ShapeUtil::ByteSizeOf(shape)); |
| 183 | auto device_buffer = std::make_shared<SharedDeviceBuffer>( |
| 184 | /*allocator=*/nullptr, local_device->device_ordinal(), |
| 185 | std::initializer_list<se::DeviceMemoryBase>{buffer}, |
| 186 | /*children=*/std::vector<std::shared_ptr<SharedDeviceBuffer>>{}, |
| 187 | /*definition_event=*/nullptr, std::move(on_delete_callback)); |
| 188 | return absl::make_unique<PyLocalBuffer>( |
| 189 | shape, shape, std::move(device_buffer), std::move(client), |
| 190 | std::move(device)); |
| 191 | } |
| 192 | |
| 193 | TransferManager* transfer_manager = |
| 194 | client->client()->backend().transfer_manager(); |
| 195 | se::DeviceMemoryAllocator* allocator = client->allocator(); |
| 196 | TF_ASSIGN_OR_RETURN(Shape compact_shape, |
| 197 | transfer_manager->ChooseCompactLayoutForShape(shape)); |
| 198 | TF_ASSIGN_OR_RETURN( |
| 199 | ScopedShapedBuffer scoped_buffer, |
| 200 | transfer_manager->AllocateScopedShapedBuffer( |
| 201 | compact_shape, allocator, local_device->device_ordinal())); |
| 202 | |
| 203 | // Make the host to device stream wait for the newly allocated buffer to be |
| 204 | // available on the compute stream. We schedule this wait synchronously; while |
| 205 | // not strictly necessary, we must not create stream dependency cycles, and |
| 206 | // adding the wait synchronously avoids any chance of any dependent |
| 207 | // computations that depend on this transfer being enqueued on the compute |
| 208 | // stream. |
| 209 | if (!transfer_manager->CanShapedBufferBeAccessedNow( |
| 210 | local_device->host_to_device_stream()->parent(), scoped_buffer)) { |
| 211 | local_device->host_to_device_stream()->ThenWaitFor( |
| 212 | local_device->compute_stream()); |
| 213 | } |
| 214 | |
| 215 | std::shared_ptr<BufferDefinitionEvent> definition_event = |
nothing calls this directly
no test coverage detected