DLPack: PyTorch leaks empty tensors

[wjakob]

🐛 Describe the bug

PyTorch supports the DLPack interface to enable zero-copy data exchange with various other array programming frameworks.

I've been investigating memory leaks that occur when exchanging tensors with PyTorch and pinned them down to an issue related to situations where torch.utils.dlpack.from_dlpack is called with an empty tensor. In this case, PyTorch fails to clean up the associated DLPack data structures, and this leads to small memory leaks. The more fundamental issue is that DLPack tensors are often backed by associated Python data structures, and this can cause numerous reference leaks there as well.

Here is a self-contained demonstrator which uses ctypes to manually create a DLPack tensor. By default, it creates a tensor containing the entries 0..5 and prints them to the screen. The script below produces the following output:

tensor([0., 1., 2., 3., 4.])
Deleter called!

We can observe how both the data was correctly handed over to PyTorch, and how it calls the cleanup routine when the tensor is no longer needed.

import ctypes as c
import torch

float_p    = c.POINTER(c.c_float)
longlong_p = c.POINTER(c.c_longlong)

class DLTensor(c.Structure):
    _fields_ = [
        ('data', float_p),
        ('device_type', c.c_int),
        ('device_id', c.c_int),
        ('ndim', c.c_int),
        ('code', c.c_uint8),
        ('bits', c.c_uint8),
        ('lanes', c.c_uint16),
        ('shape', longlong_p),
        ('strides', longlong_p),
        ('byte_offset', c.c_ulonglong),
        ('ctx', c.c_void_p),
        ('deleter', c.CFUNCTYPE(None, c.c_void_p))
    ]    

@c.CFUNCTYPE(None, c.c_void_p)
def deleter(tensor):
    print('Deleter called!')


def make_tensor(size):
    tensor = DLTensor()

    if size != 0:
        data = (c.c_float * size)()
        for i in range(size):
            data[i] = c.c_float(i)
    else:
        data = None

    shape = (c.c_longlong * 1)()
    shape[0] = size

    strides = (c.c_longlong * 1)()
    strides[0] = 1

    tensor.data = data
    tensor.device_type = 1 # cpu
    tensor.device_id = 0
    tensor.ndim = 1
    tensor.code = 2 # float
    tensor.lanes = 1
    tensor.bits = 32
    tensor.shape = shape
    tensor.strides = strides
    tensor.ctx = 0
    tensor.deleter = deleter


    c.pythonapi.PyCapsule_New.argtypes = [c.c_void_p, c.c_char_p, c.c_void_p]
    c.pythonapi.PyCapsule_New.restype = c.py_object

    capsule = c.pythonapi.PyCapsule_New(c.cast(c.pointer(tensor), c.c_void_p), b'dltensor', None)

    a = torch.utils.dlpack.from_dlpack(capsule)
    print(a)


make_tensor(5)

Now, change the last line to make_tensor(0). In this case, the output is

tensor([])

The deleter is no longer called, presumably because there is nothing to cleanup. But this is not correct. The DLPack specification mandates that a framework which consumes a tensor must eventually call its delete member. And we can see that PyTorch consumed this tensor by adding a print(capsule) statement at the end of the function make_tensor().

This produces the output

<capsule object "used_dltensor" at 0x1320eafa0>

Versions

PyTorch version: 2.3.0.dev20240108
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.6.2 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.28.1
Libc version: N/A

Python version: 3.12.1 (main, Dec 8 2023, 18:57:37) [Clang 14.0.3 (clang-1403.0.22.14.1)] (64-bit runtime)
Python platform: macOS-13.6.2-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Max

Versions of relevant libraries:
[pip3] numpy==1.26.2
[pip3] torch==2.3.0.dev20240108
[conda] Could not collect

[malfet]

Isn't it the same as #117058 ?

[wjakob]

Yes, that looks like exactly the same issue.

[albanD]

Closing as a duplicate of #117058