| 289 | } |
| 290 | |
| 291 | StatusOr<std::unique_ptr<PyLocalBuffer>> DLPackManagedTensorToBuffer( |
| 292 | const pybind11::capsule& tensor, std::shared_ptr<PyLocalClient> client) { |
| 293 | if (absl::string_view(tensor.name()) != kDlTensorCapsuleName) { |
| 294 | return InvalidArgument( |
| 295 | "DLPack tensor must be a capsule with name \"dltensor\", got \"%s\". " |
| 296 | "Note that a DLPack tensor may be consumed at most once.", |
| 297 | absl::string_view(tensor.name())); |
| 298 | } |
| 299 | DLManagedTensor* dlmt = static_cast<DLManagedTensor*>(tensor); |
| 300 | if (dlmt->dl_tensor.ndim < 0) { |
| 301 | return InvalidArgument( |
| 302 | "Number of dimensions in DLManagedTensor must be nonnegative, got %d", |
| 303 | dlmt->dl_tensor.ndim); |
| 304 | } |
| 305 | TF_ASSIGN_OR_RETURN(std::shared_ptr<Device> device, |
| 306 | DeviceForDLContext(*client, dlmt->dl_tensor.ctx)); |
| 307 | absl::Span<int64 const> dimensions( |
| 308 | reinterpret_cast<int64*>(dlmt->dl_tensor.shape), dlmt->dl_tensor.ndim); |
| 309 | TF_ASSIGN_OR_RETURN(PrimitiveType element_type, |
| 310 | DLDataTypeToPrimitiveType(dlmt->dl_tensor.dtype)); |
| 311 | |
| 312 | std::vector<int64> minor_to_major; |
| 313 | if (dlmt->dl_tensor.strides && !absl::c_find(dimensions, 0)) { |
| 314 | absl::Span<int64 const> strides( |
| 315 | reinterpret_cast<int64*>(dlmt->dl_tensor.strides), |
| 316 | dlmt->dl_tensor.ndim); |
| 317 | TF_ASSIGN_OR_RETURN(minor_to_major, StridesToLayout(dimensions, strides)); |
| 318 | } else { |
| 319 | minor_to_major.resize(dlmt->dl_tensor.ndim); |
| 320 | std::iota(minor_to_major.rbegin(), minor_to_major.rend(), 0); |
| 321 | } |
| 322 | Shape shape = |
| 323 | ShapeUtil::MakeShapeWithLayout(element_type, dimensions, minor_to_major); |
| 324 | se::DeviceMemoryBase buffer( |
| 325 | static_cast<char*>(dlmt->dl_tensor.data) + dlmt->dl_tensor.byte_offset, |
| 326 | ShapeUtil::ByteSizeOf(shape)); |
| 327 | |
| 328 | std::function<void()> on_delete_callback; |
| 329 | if (dlmt->deleter) { |
| 330 | on_delete_callback = [dlmt]() { dlmt->deleter(dlmt); }; |
| 331 | } |
| 332 | auto device_buffer = std::make_shared<SharedDeviceBuffer>( |
| 333 | /*allocator=*/nullptr, dlmt->dl_tensor.ctx.device_id, |
| 334 | std::initializer_list<se::DeviceMemoryBase>{buffer}, |
| 335 | /*children=*/std::vector<std::shared_ptr<SharedDeviceBuffer>>{}, |
| 336 | /*definition_event=*/nullptr, std::move(on_delete_callback)); |
| 337 | |
| 338 | // We have taken ownership of the array inside the capsule; make sure the |
| 339 | // capsule it cannot be used again. |
| 340 | PyCapsule_SetName(tensor.ptr(), "used_dltensor"); |
| 341 | PyCapsule_SetDestructor(tensor.ptr(), nullptr); |
| 342 | return absl::make_unique<PyLocalBuffer>(shape, shape, |
| 343 | std::move(device_buffer), |
| 344 | std::move(client), std::move(device)); |
| 345 | } |
| 346 | |
| 347 | } // namespace xla |
nothing calls this directly
no test coverage detected