| 555 | |
| 556 | template <typename Backend> |
| 557 | DLMTensorPtr ToDLMTensor(Tensor<Backend> &tensor, |
| 558 | std::optional<intptr_t> stream_handle_value, |
| 559 | std::optional<std::pair<DLDeviceType, int>> dl_device) { |
| 560 | DLDevice dev; |
| 561 | |
| 562 | if (dl_device.has_value()) |
| 563 | dev = { dl_device->first, dl_device->second }; |
| 564 | else |
| 565 | dev = GetDLDevice(tensor); |
| 566 | |
| 567 | if (dev.device_type == kDLCUDA) { |
| 568 | AccessOrder target_order = cudaStreamLegacy; |
| 569 | if (stream_handle_value.has_value()) { |
| 570 | if (*stream_handle_value == -1) { |
| 571 | target_order = AccessOrder{}; |
| 572 | } else { |
| 573 | cudaStream_t stream = cudaStream_t(*stream_handle_value); |
| 574 | if (stream == 0) |
| 575 | throw std::invalid_argument("Stream 0 is explicitly forbidden by DLPack protocol"); |
| 576 | target_order = AccessOrder(stream, dev.device_id); |
| 577 | } |
| 578 | } |
| 579 | |
| 580 | if constexpr (std::is_same_v<Backend, CPUBackend>) { |
| 581 | throw std::runtime_error( |
| 582 | "The tensor is in CPU memory and a CUDA DLPack tensor was requested"); |
| 583 | } |
| 584 | if (dev.device_id != tensor.device_id()) |
| 585 | throw std::runtime_error(make_string("Requested a DLPack tensor for GPU_", dev.device_id, |
| 586 | "while the tensor resides in GPU_", tensor.device_id(), " memory.")); |
| 587 | |
| 588 | if (target_order) { |
| 589 | if (tensor.ready_event()) { |
| 590 | target_order.wait(tensor.ready_event()); |
| 591 | } else { |
| 592 | target_order.wait(tensor.order()); |
| 593 | } |
| 594 | } |
| 595 | return GetSharedDLTensor(tensor); |
| 596 | } else if (dev.device_type == kDLCPU || dev.device_type == kDLCUDAHost) { |
| 597 | if constexpr (std::is_same_v<Backend, GPUBackend>) { |
| 598 | throw std::runtime_error( |
| 599 | "The tensor is in CUDA GPU memory and a CPU DLPack tensor was requested"); |
| 600 | } |
| 601 | |
| 602 | if (dev.device_type == kDLCUDAHost && !tensor.is_pinned()) |
| 603 | throw std::runtime_error( |
| 604 | "A CUDA host (pinned) DLPack was requested, but the tensor buffer is not pinned."); |
| 605 | |
| 606 | if (tensor.is_pinned() && tensor.ready_event()) { |
| 607 | // DLPack doesn't support stream-ordered CUDA host tensors |
| 608 | AccessOrder::host().wait(tensor.ready_event()); |
| 609 | } |
| 610 | |
| 611 | DLMTensorPtr dlm_tensor = GetSharedDLTensor(tensor); |
| 612 | // Set the device type to the desired one - if the original tensor was pinned, we can |
| 613 | // downgrade it to regular host memory. |
| 614 | dlm_tensor->dl_tensor.device = dev; |
no test coverage detected