Spaces:
Sleeping
Sleeping
| from typing import Any | |
| import torch | |
| import enum | |
| from torch._C import _from_dlpack | |
| from torch._C import _to_dlpack as to_dlpack | |
| class DLDeviceType(enum.IntEnum): | |
| # Enums as in DLPack specification (aten/src/ATen/dlpack.h) | |
| kDLCPU = 1, | |
| kDLGPU = 2, | |
| kDLCPUPinned = 3, | |
| kDLOpenCL = 4, | |
| kDLVulkan = 7, | |
| kDLMetal = 8, | |
| kDLVPI = 9, | |
| kDLROCM = 10, | |
| kDLExtDev = 12, | |
| kDLOneAPI = 14, | |
| torch._C._add_docstr(to_dlpack, r"""to_dlpack(tensor) -> PyCapsule | |
| Returns an opaque object (a "DLPack capsule") representing the tensor. | |
| .. note:: | |
| ``to_dlpack`` is a legacy DLPack interface. The capsule it returns | |
| cannot be used for anything in Python other than use it as input to | |
| ``from_dlpack``. The more idiomatic use of DLPack is to call | |
| ``from_dlpack`` directly on the tensor object - this works when that | |
| object has a ``__dlpack__`` method, which PyTorch and most other | |
| libraries indeed have now. | |
| .. warning:: | |
| Only call ``from_dlpack`` once per capsule produced with ``to_dlpack``. | |
| Behavior when a capsule is consumed multiple times is undefined. | |
| Args: | |
| tensor: a tensor to be exported | |
| The DLPack capsule shares the tensor's memory. | |
| """) | |
| # TODO: add a typing.Protocol to be able to tell Mypy that only objects with | |
| # __dlpack__ and __dlpack_device__ methods are accepted. | |
| def from_dlpack(ext_tensor: Any) -> 'torch.Tensor': | |
| """from_dlpack(ext_tensor) -> Tensor | |
| Converts a tensor from an external library into a ``torch.Tensor``. | |
| The returned PyTorch tensor will share the memory with the input tensor | |
| (which may have come from another library). Note that in-place operations | |
| will therefore also affect the data of the input tensor. This may lead to | |
| unexpected issues (e.g., other libraries may have read-only flags or | |
| immutable data structures), so the user should only do this if they know | |
| for sure that this is fine. | |
| Args: | |
| ext_tensor (object with ``__dlpack__`` attribute, or a DLPack capsule): | |
| The tensor or DLPack capsule to convert. | |
| If ``ext_tensor`` is a tensor (or ndarray) object, it must support | |
| the ``__dlpack__`` protocol (i.e., have a ``ext_tensor.__dlpack__`` | |
| method). Otherwise ``ext_tensor`` may be a DLPack capsule, which is | |
| an opaque ``PyCapsule`` instance, typically produced by a | |
| ``to_dlpack`` function or method. | |
| Examples:: | |
| >>> import torch.utils.dlpack | |
| >>> t = torch.arange(4) | |
| # Convert a tensor directly (supported in PyTorch >= 1.10) | |
| >>> t2 = torch.from_dlpack(t) | |
| >>> t2[:2] = -1 # show that memory is shared | |
| >>> t2 | |
| tensor([-1, -1, 2, 3]) | |
| >>> t | |
| tensor([-1, -1, 2, 3]) | |
| # The old-style DLPack usage, with an intermediate capsule object | |
| >>> capsule = torch.utils.dlpack.to_dlpack(t) | |
| >>> capsule | |
| <capsule object "dltensor" at ...> | |
| >>> t3 = torch.from_dlpack(capsule) | |
| >>> t3 | |
| tensor([-1, -1, 2, 3]) | |
| >>> t3[0] = -9 # now we're sharing memory between 3 tensors | |
| >>> t3 | |
| tensor([-9, -1, 2, 3]) | |
| >>> t2 | |
| tensor([-9, -1, 2, 3]) | |
| >>> t | |
| tensor([-9, -1, 2, 3]) | |
| """ | |
| if hasattr(ext_tensor, '__dlpack__'): | |
| device = ext_tensor.__dlpack_device__() | |
| # device is either CUDA or ROCm, we need to pass the current | |
| # stream | |
| if device[0] in (DLDeviceType.kDLGPU, DLDeviceType.kDLROCM): | |
| stream = torch.cuda.current_stream(f'cuda:{device[1]}') | |
| # cuda_stream is the pointer to the stream and it is a public | |
| # attribute, but it is not documented | |
| # The array API specify that the default legacy stream must be passed | |
| # with a value of 1 for CUDA | |
| # https://data-apis.org/array-api/latest/API_specification/array_object.html?dlpack-self-stream-none#dlpack-self-stream-none | |
| is_cuda = device[0] == DLDeviceType.kDLGPU | |
| # Since pytorch is not using PTDS by default, lets directly pass | |
| # the legacy stream | |
| stream_ptr = 1 if is_cuda and stream.cuda_stream == 0 else stream.cuda_stream | |
| dlpack = ext_tensor.__dlpack__(stream=stream_ptr) | |
| else: | |
| dlpack = ext_tensor.__dlpack__() | |
| else: | |
| # Old versions just call the converter | |
| dlpack = ext_tensor | |
| return _from_dlpack(dlpack) | |