Repository: cupy/cupy

Issue:
Make DLPack a submodule?
I am looking into DLPack as we likely need to fix it for ROCm (it supports the `kDLROCM` option). One thing I noticed is that we basically reimplemented everything from the DLPack header [`dlpack.h`](https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h) in `cupy/core/dlpack.pyx`. I suppose like Jitify it's also a single header thing, so we might as well make it a submodule? 

This will also prepare us for the Python Array API involvement, for which DLPack is being considered as the data exchange format (https://github.com/data-apis/consortium-feedback/issues/1; see also the [Array API Standard draft](https://data-apis.github.io/array-api/latest/design_topics/data_interchange.html)), but it currently lacks 
1. support for complex numbers (https://github.com/dmlc/dlpack/issues/50)
2. specification for synchronization semantics as in CUDA Array Interface v3 (https://github.com/dmlc/dlpack/issues/57)

so it will likely be revised in the near future. If this is acceptable then please assign it to me.


Discussion:

Resolving patch:
diff --git a/cupy_setup_build.py b/cupy_setup_build.py
--- a/cupy_setup_build.py
+++ b/cupy_setup_build.py
@@ -68,7 +68,6 @@
     'cupy.core._routines_statistics',
     'cupy.core._scalar',
     'cupy.core.core',
-    'cupy.core.dlpack',
     'cupy.core.flags',
     'cupy.core.internal',
     'cupy.core.fusion',
@@ -323,6 +322,17 @@
             'version_method': build.get_thrust_version,
         })
 
+MODULES.append({
+    'name': 'dlpack',
+    'file': [
+        'cupy.core.dlpack',
+    ],
+    'include': [
+        'cupy/dlpack/dlpack.h',
+    ],
+    'libraries': [],
+})
+
 
 def ensure_module_file(file):
     if isinstance(file, tuple):
@@ -605,6 +615,10 @@ def make_extensions(options, compiler, use_cython):
             # if any change is made to the Jitify header, we force recompiling
             s['depends'] = ['./cupy/core/include/cupy/jitify/jitify.hpp']
 
+        if module['name'] == 'dlpack':
+            # if any change is made to the DLPack header, we force recompiling
+            s['depends'] = ['./cupy/core/include/cupy/dlpack/dlpack.h']
+
         for f in module['file']:
             s_file = copy.deepcopy(s)
             name = module_extension_name(f)


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_dlpack.py b/tests/cupy_tests/core_tests/test_dlpack.py
--- a/tests/cupy_tests/core_tests/test_dlpack.py
+++ b/tests/cupy_tests/core_tests/test_dlpack.py
@@ -1,5 +1,7 @@
 import unittest
 
+import pytest
+
 import cupy
 from cupy import testing
 
@@ -38,13 +40,41 @@ def setUp(self):
         cupy.cuda.set_allocator(self.pool.malloc)
 
     def tearDown(self):
+        self.pool.free_all_blocks()
         cupy.cuda.set_allocator(self.old_pool.malloc)
 
     def test_deleter(self):
+        # memory is freed when tensor is deleted, as it's not consumed
         array = cupy.empty(10)
         tensor = array.toDlpack()
+        # str(tensor): <capsule object "dltensor" at 0x7f7c4c835330>
+        assert "\"dltensor\"" in str(tensor)
         assert self.pool.n_free_blocks() == 0
         del array
         assert self.pool.n_free_blocks() == 0
         del tensor
         assert self.pool.n_free_blocks() == 1
+
+    def test_deleter2(self):
+        # memory is freed when array2 is deleted, as tensor is consumed
+        array = cupy.empty(10)
+        tensor = array.toDlpack()
+        assert "\"dltensor\"" in str(tensor)
+        array2 = cupy.fromDlpack(tensor)
+        assert "\"used_dltensor\"" in str(tensor)
+        assert self.pool.n_free_blocks() == 0
+        del array
+        assert self.pool.n_free_blocks() == 0
+        del array2
+        assert self.pool.n_free_blocks() == 1
+        del tensor
+        assert self.pool.n_free_blocks() == 1
+
+    def test_multiple_consumption_error(self):
+        # Prevent segfault, see #3611
+        array = cupy.empty(10)
+        tensor = array.toDlpack()
+        array2 = cupy.fromDlpack(tensor)  # noqa
+        with pytest.raises(ValueError) as e:
+            array3 = cupy.fromDlpack(tensor)  # noqa
+        assert 'consumed multiple times' in str(e.value)