Refactor. Add tests for negative strides and reduce test cases.

python/pyarrow/_cuda.pyx

@@ -15,19 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import functools
 
 from pyarrow.compat import tobytes
 from pyarrow.lib cimport *
 from pyarrow.includes.libarrow_cuda cimport *
 from pyarrow.lib import py_buffer, allocate_buffer, as_buffer
+from pyarrow.util import get_contiguous_span
 cimport cpython as cp
 
 
-def product(seq):
-    return functools.reduce(lambda a, b: a*b, seq, 1)
-
-
 cdef class Context:
     """
     CUDA driver context.
@@ -198,7 +194,8 @@ cdef class Context:
         """Create device buffer from address and size as a view.
 
         The caller is responsible for allocating and freeing the
-        memory.
+        memory. When `address==size==0` then a new zero-sized buffer
+        is returned.
 
         Parameters
         ----------
@@ -214,7 +211,10 @@ cdef class Context:
         -------
         cbuf : CudaBuffer
           Device buffer as a view of device reachable memory.
+
         """
+        if not address and size == 0:
+            return self.new_buffer(0)
         cdef:
             uintptr_t c_addr = self.get_device_address(address)
             int64_t c_size = size
@@ -309,46 +309,24 @@ cdef class Context:
 
         """
         if isinstance(obj, HostBuffer):
-            if not obj.address and obj.size == 0:
-                return self.new_buffer(0)
             return self.foreign_buffer(obj.address, obj.size)
         elif isinstance(obj, Buffer):
             return CudaBuffer.from_buffer(obj)
         elif isinstance(obj, CudaBuffer):
             return obj
         elif hasattr(obj, '__cuda_array_interface__'):
-            import numpy as np
             desc = obj.__cuda_array_interface__
-            shape = desc['shape']
-            strides = desc.get('strides')
-            dtype = np.dtype(desc['typestr'])
             addr = desc['data'][0]
             if addr is None:
-                nbytes = 0
-            elif not strides:
-                nbytes = dtype.itemsize * product(shape)
-            else:
-                start = 0
-                end = dtype.itemsize
-                for i, dim in enumerate(shape):
-                    if dim == 0:
-                        start = end = addr
-                        break
-                    stride = strides[i]
-                    if stride > 0:
-                        end += stride * (dim - 1)
-                    elif stride < 0:
-                        start += stride * (dim - 1)
-                nbytes = end - start
-            if nbytes:
-                return self.foreign_buffer(addr + start, nbytes)
-            else:
                 return self.new_buffer(0)
+            import numpy as np
+            start, end = get_contiguous_span(
+                desc['shape'], desc.get('strides'),
+                np.dtype(desc['typestr']).itemsize)
+            return self.foreign_buffer(addr + start, end - start)
         from numba.cuda.cudadrv.driver import MemoryPointer
         if isinstance(obj, MemoryPointer):
             addr = obj.device_pointer.value
-            if not addr and obj.size == 0:
-                return self.new_buffer(0)
             return self.foreign_buffer(addr, obj.size)
         raise NotImplementedError('cannot create device buffer view from'
                                   ' `%s` object' % (type(obj)))

python/pyarrow/tests/test_cuda.py

@@ -60,7 +60,7 @@ def test_Context():
         cuda.Context(cuda.Context.get_num_devices())
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_manage_allocate_free_host(size):
     buf = cuda.new_host_buffer(size)
     arr = np.frombuffer(buf, dtype=np.uint8)
@@ -102,7 +102,7 @@ def make_random_buffer(size, target='host'):
     raise ValueError('invalid target value')
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_context_device_buffer(size):
     # Creating device buffer from host buffer;
     arr, buf = make_random_buffer(size)
@@ -262,7 +262,7 @@ def test_context_from_object(size):
         ctx.buffer_from_object(np.array([1, 2, 3]))
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_CudaBuffer(size):
     arr, buf = make_random_buffer(size)
     assert arr.tobytes() == buf.to_pybytes()
@@ -287,7 +287,7 @@ def test_CudaBuffer(size):
         cuda.CudaBuffer()
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_HostBuffer(size):
     arr, buf = make_random_buffer(size)
     assert arr.tobytes() == buf.to_pybytes()
@@ -313,7 +313,7 @@ def test_HostBuffer(size):
         cuda.HostBuffer()
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_copy_from_to_host(size):
 
     # Create a buffer in host containing range(size)
@@ -338,7 +338,7 @@ def test_copy_from_to_host(size):
     np.testing.assert_equal(arr, arr2)
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_copy_to_host(size):
     arr, dbuf = make_random_buffer(size, target='device')
 
@@ -398,7 +398,7 @@ def test_copy_to_host(size):
             dbuf.copy_to_host(buf=buf, position=position, nbytes=nbytes)
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_copy_from_device(size):
     arr, buf = make_random_buffer(size=size, target='device')
     lst = arr.tolist()
@@ -442,7 +442,7 @@ def test_copy_from_device(size):
             put(position=position, nbytes=nbytes)
 
 
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_copy_from_host(size):
     arr, buf = make_random_buffer(size=size, target='host')
     lst = arr.tolist()
@@ -649,7 +649,7 @@ def other_process_for_test_IPC(handle_buffer, expected_arr):
 
 @cuda_ipc
 @pytest.mark.skipif(sys.version_info[0] == 2, reason="test needs Python 3")
-@pytest.mark.parametrize("size", [0, 1, 8, 1000])
+@pytest.mark.parametrize("size", [0, 1, 1000])
 def test_IPC(size):
     import multiprocessing
     ctx = multiprocessing.get_context('spawn')

python/pyarrow/tests/test_cuda_numba_interop.py

@@ -78,7 +78,7 @@ def make_random_buffer(size, target='host', dtype='uint8', ctx=None):
 @pytest.mark.parametrize("c", range(len(context_choice_ids)),
                          ids=context_choice_ids)
 @pytest.mark.parametrize("dtype", dtypes, ids=dtypes)
-@pytest.mark.parametrize("size", [0, 1, 1000])
+@pytest.mark.parametrize("size", [0, 1, 8, 1000])
 def test_from_object(c, dtype, size):
     ctx, nb_ctx = context_choices[c]
     arr, cbuf = make_random_buffer(size, target='device', dtype=dtype, ctx=ctx)
@@ -91,7 +91,10 @@ def test_from_object(c, dtype, size):
     np.testing.assert_equal(arr, arr2)
 
     # Creating device buffer from numba DeviceNDArray:
-    darr = DeviceNDArray((size,), (np.dtype(dtype).itemsize,), np.dtype(dtype),
+    darr = nb_cuda.to_device(arr)
+    darr = DeviceNDArray(size,
+                         np.dtype(dtype).itemsize,
+                         np.dtype(dtype),
                          gpu_data=mem)
     cbuf2 = ctx.buffer_from_object(darr)
     assert cbuf2.size == cbuf.size
@@ -101,21 +104,36 @@ def test_from_object(c, dtype, size):
     # Creating device buffer from a slice of numba DeviceNDArray:
     if size >= 8:
         for s in [slice(size//4, None, None),
-                  slice(None, -(size//4), None),
                   slice(size//4, -(size//4), None),
                   slice(None, None, 2),
                   slice(size//4, None, 2),
-                  slice(None, -(size//4), 2),
                   slice(size//4, -(size//4), 2),
                   slice(None, None, 3),
                   slice(size//8, None, 3),
-                  slice(None, -(size//8), 3),
                   slice(size//8, -(size//8), 3)]:
             s2 = slice(None, None, s.step)
             cbuf2 = ctx.buffer_from_object(darr[s])
             arr2 = np.frombuffer(cbuf2.copy_to_host(), dtype=dtype)
             np.testing.assert_equal(arr[s], arr2[s2])
 
+        # slice with negative strides
+        rdarr = darr[::-1]
+        if 1:
+            # workaround numba bug [numba issue 3705]:
+            from numba.cuda.cudadrv.driver import MemoryPointer
+            import ctypes
+            rdarr.shape = arr[::-1].shape
+            addr = (darr.__cuda_array_interface__['data'][0]
+                    + (arr[::-1].__array_interface__['data'][0]
+                       - arr.__array_interface__['data'][0]))
+            mem = MemoryPointer(nb_ctx, pointer=ctypes.c_void_p(addr),
+                                size=size)
+            rdarr.gpu_data = mem
+        cbuf2 = ctx.buffer_from_object(rdarr)
+        assert cbuf2.size == cbuf.size
+        arr2 = np.frombuffer(cbuf2.copy_to_host(), dtype=dtype)
+        np.testing.assert_equal(arr, arr2)
+
     # Creating device buffer from a 2-dimensional numba DeviceNDArray:
     if size >= 8:
         cbuf2 = ctx.buffer_from_object(darr.reshape(size//4, size//(size//4)))

python/pyarrow/util.py

@@ -19,6 +19,8 @@
 
 import six
 import warnings
+import functools
+
 
 try:
     from textwrap import indent
@@ -78,3 +80,44 @@ def _stringify_path(path):
             return str(path)
 
     raise TypeError("not a path-like object")
+
+
+def product(seq):
+    """
+    Return a product of sequence items.
+    """
+    return functools.reduce(lambda a, b: a*b, seq, 1)
+
+
+def get_contiguous_span(shape, strides, itemsize):
+    """
+    Return a span of N-D array data.
+
+    Parameters
+    ----------
+    shape : tuple
+    strides : tuple
+    itemsize : int
+      Specify array shape data
+
+    Returns
+    -------
+    start, end : int
+      The span end points.
+    """
+    if not strides:
+        start = 0
+        end = itemsize * product(shape)
+    else:
+        start = 0
+        end = itemsize
+        for i, dim in enumerate(shape):
+            if dim == 0:
+                start = end = 0
+                break
+            stride = strides[i]
+            if stride > 0:
+                end += stride * (dim - 1)
+            elif stride < 0:
+                start += stride * (dim - 1)
+    return start, end