如何從gpu內存地址創建PyCUDA GPUArray？

Question

我正在使用PyTorch，並希望在PyCUDA的幫助下對Tensor數據進行一些算術運算。 我可以得到一個CUDA張量的內存地址t通過t.data_ptr() 我可以以某種方式使用此地址以及我對大小和數據類型的了解來初始化GPUArray嗎？ 我希望避免復制數據，但這也是一種替代方案。

Answer 1

事實證明這是可能的。 我們需要一個指針來處理數據，這需要一些額外的功能：

class Holder(PointerHolderBase):

    def __init__(self, tensor):
        super().__init__()
        self.tensor = tensor
        self.gpudata = tensor.data_ptr()

    def get_pointer(self):
        return self.tensor.data_ptr()

    def __int__(self):
        return self.__index__()

    # without an __index__ method, arithmetic calls to the GPUArray backed by this pointer fail
    # not sure why, this needs to return some integer, apparently
    def __index__(self):
        return self.gpudata

然后我們可以使用這個類來實例化GPUArray 。 代碼使用Reikna數組，它是一個子類，但也應該與pycuda數組一起使用。

def tensor_to_gpuarray(tensor, context=pycuda.autoinit.context):
    '''Convert a :class:`torch.Tensor` to a :class:`pycuda.gpuarray.GPUArray`. The underlying
    storage will be shared, so that modifications to the array will reflect in the tensor object.
    Parameters
    ----------
    tensor  :   torch.Tensor
    Returns
    -------
    pycuda.gpuarray.GPUArray
    Raises
    ------
    ValueError
        If the ``tensor`` does not live on the gpu
    '''
    if not tensor.is_cuda:
        raise ValueError('Cannot convert CPU tensor to GPUArray (call `cuda()` on it)')
    else:
        thread = cuda.cuda_api().Thread(context)
    return reikna.cluda.cuda.Array(thread, tensor.shape, dtype=torch_dtype_to_numpy(tensor.dtype), base_data=Holder(tensor))

我們可以回過頭來看看這段代碼。 我沒有找到一種方法來做到這一點而不復制數據。

def gpuarray_to_tensor(gpuarray, context=pycuda.autoinit.context):
    '''Convert a :class:`pycuda.gpuarray.GPUArray` to a :class:`torch.Tensor`. The underlying
    storage will NOT be shared, since a new copy must be allocated.
    Parameters
    ----------
    gpuarray  :   pycuda.gpuarray.GPUArray
    Returns
    -------
    torch.Tensor
    '''
    shape = gpuarray.shape
    dtype = gpuarray.dtype
    out_dtype = numpy_dtype_to_torch(dtype)
    out = torch.zeros(shape, dtype=out_dtype).cuda()
    gpuarray_copy = tensor_to_gpuarray(out, context=context)
    byte_size = gpuarray.itemsize * gpuarray.size
    pycuda.driver.memcpy_dtod(gpuarray_copy.gpudata, gpuarray.gpudata, byte_size)
    return out

---

老答案

from pycuda.gpuarray import GPUArray


def torch_dtype_to_numpy(dtype):
    dtype_name = str(dtype)[6:]     # remove 'torch.'
    return getattr(np, dtype_name)


def tensor_to_gpuarray(tensor):
    if not tensor.is_cuda:
        raise ValueError('Cannot convert CPU tensor to GPUArray (call `cuda()` on it)')
    else:
        array = GPUArray(tensor.shape, dtype=torch_dtype_to_numpy(tensor.dtype),
                         gpudata=tensor.data_ptr())
        return array.copy()

不幸的是，傳遞一個int作為gpudata關鍵字（或亞型pycuda.driver.PointerHolderBase作為pytorch論壇建議）似乎表面上工作，但許多操作失敗，看似無關的錯誤。 復制數組似乎將其轉換為可用的格式。 我認為這與gpudata成員應該是pycuda.driver.DeviceAllocation對象這一事實有關，它似乎無法從Python實例化。

現在，如何從原始數據返回到Tensor是另一回事。