torch.utils.dlpack 的源代码

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):
    # 枚举值如DLPack规范中所定义（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

返回一个表示张量的不透明对象（一个“DLPack胶囊”）。

.. 注意::
  ``to_dlpack`` 是一个遗留的DLPack接口。它返回的胶囊不能在Python中用于其他任何用途，只能作为输入传递给``from_dlpack``。更惯用的DLPack使用方法是直接在张量对象上调用``from_dlpack``——当该对象具有``__dlpack__``方法时，这将会起作用，PyTorch和大多数其他库现在确实都有这个方法。

.. 警告::
  每个由``to_dlpack``生成的胶囊只能调用一次``from_dlpack``。多次消费胶囊的行为是未定义的。

参数:
    tensor: 要导出的张量

DLPack胶囊与张量的内存共享。
""")


# TODO: 添加一个typing.Protocol，以便能够告诉Mypy只接受具有__dlpack__和__dlpack_device__方法的对象。
def from_dlpack(ext_tensor: Any) -> 'torch.Tensor':
    """from_dlpack(ext_tensor) -> Tensor

    将来自外部库的张量转换为``torch.Tensor``。

    返回的PyTorch张量将与输入张量共享内存（该张量可能来自另一个库）。请注意，就地操作也将影响输入张量的数据。这可能会导致意外问题（例如，其他库可能有只读标志或不可变的数据结构），因此用户只有在确定这样做没问题时才应执行此操作。

    参数:
        ext_tensor (具有``__dlpack__``属性的对象，或DLPack胶囊):
            要转换的张量或DLPack胶囊。

            如果``ext_tensor``是一个张量（或ndarray）对象，它必须支持``__dlpack__``协议（即，具有``ext_tensor.__dlpack__``方法）。否则，``ext_tensor``可能是一个DLPack胶囊，这是一个不透明的``PyCapsule``实例，通常由``to_dlpack``函数或方法生成。

    示例::

        >>> import torch.utils.dlpack
        >>> t = torch.arange(4)

        # 直接转换张量（支持在PyTorch >= 1.10中）
        >>> t2 = torch.from_dlpack(t)
        >>> t2[:2] = -1  # 显示内存是共享的
        >>> t2
        tensor([-1, -1,  2,  3])
        >>> t
        tensor([-1, -1,  2,  3])

        # 旧式的DLPack用法，带有中间的胶囊对象
        >>> capsule = torch.utils.dlpack.to_dlpack(t)
        >>> capsule
        
        >>> t3 = torch.from_dlpack(capsule)
        >>> t3
        tensor([-1, -1,  2,  3])
        >>> t3[0] = -9  # 现在我们在3个张量之间共享内存
        >>> 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 要么是CUDA要么是ROCm，我们需要传递当前的流
        if device[0] in (DLDeviceType.kDLGPU, DLDeviceType.kDLROCM):
            stream = torch.cuda.current_stream(f'cuda:{device[1]}')
            # cuda_stream 是流的指针，它是一个公共属性，但没有文档说明
            # 数组API指定必须传递默认的遗留流，对于CUDA，值为1
            # 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
            # 由于pytorch默认不使用PTDS，让我们直接传递遗留流
            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:
        # 旧版本直接调用转换器
        dlpack = ext_tensor
    return _from_dlpack(dlpack)